diff --git a/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/DownloadTests.cs b/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/DownloadTests.cs index aea0b91b..48690f87 100644 --- a/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/DownloadTests.cs +++ b/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/DownloadTests.cs @@ -1,7 +1,12 @@ -using System; +using Agent.Utilities; +using SshNet.Security.Cryptography; +using System; using System.Collections.Generic; using System.Linq; +using System.Net; +using System.Security.Cryptography; using System.Text; +using System.Text.Json; using System.Threading.Tasks; namespace Agent.Tests.PluginTests @@ -17,11 +22,11 @@ public class DownloadTests ICryptoManager _cryptoManager = new TestCryptoManager(); IMessageManager _messageManager = new TestMessageManager(); ISpawner _spawner = new TestSpawner(); - IPlugin _downloadPlugin { get; set; } + IFilePlugin _downloadPlugin { get; set; } ServerJob _downloadJob { get; set; } public DownloadTests() { - _downloadPlugin = PluginLoader.LoadPluginFromDisk("download", _messageManager, _config, _logger, _tokenManager, _spawner); + _downloadPlugin = (IFilePlugin)PluginLoader.LoadPluginFromDisk("download", _messageManager, _config, _logger, _tokenManager, _spawner); _downloadJob = new ServerJob() { task = new ServerTask() @@ -78,15 +83,138 @@ public void TestMultiChunkDownload() public void TestSingleChunkDownload() { string fileName = Path.GetTempPath() + Guid.NewGuid().ToString() + ".txt"; - Utilities.CreateTemporaryFileWithRandomText(fileName, 512000 * 3); + Utilities.CreateTemporaryFileWithRandomText(fileName, 512000); + _config.chunk_size = 512000; Dictionary downloadParams = new Dictionary() { + { "host", Dns.GetHostName() }, { "path", fileName }, }; - Assert.IsTrue(false); + _downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams); + _downloadPlugin.Execute(_downloadJob); + + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + DownloadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + ServerResponseResult responseResult = new ServerResponseResult() + { + task_id = "123", + file_id = "1234", + total_chunks = 4, + chunk_num = 1, + status = "success" + }; + _downloadPlugin.HandleNextMessage(responseResult); + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + + Assert.IsNotNull(ur.download.chunk_data); + Assert.AreNotEqual(Misc.Base64DecodeToByteArray(ur.download.chunk_data).Length, 0); + Assert.AreEqual(GetHashForByteArray(Misc.Base64DecodeToByteArray(ur.download.chunk_data)), GetHashForFile(fileName)); + } + [TestMethod] + public void TestHandleNextChunkFailure() + { + string fileName = Path.GetTempPath() + Guid.NewGuid().ToString() + ".txt"; + Utilities.CreateTemporaryFileWithRandomText(fileName, 512000); + _config.chunk_size = 512000; + Dictionary downloadParams = new Dictionary() + { + { "host", Dns.GetHostName() }, + { "path", fileName }, + + }; + _downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams); + _downloadPlugin.Execute(_downloadJob); + + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + DownloadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + ServerResponseResult responseResult = new ServerResponseResult() + { + task_id = "123", + file_id = "1234", + total_chunks = 4, + chunk_num = 1, + status = "failed" + }; + _downloadPlugin.HandleNextMessage(responseResult); + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + + Assert.AreEqual(ur.status, "error"); //Test to make sure the plugin parses local paths like we expect } + + [TestMethod] + public void TestUncPathParsing() + { + string hostName = "127.0.0.1"; + string filePath = "C$\\Windows\\System32\\drivers\\etc\\hosts"; + Dictionary downloadParams = new Dictionary() + { + {"host", hostName }, + {"path", filePath }, + + }; + _downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams); + _downloadPlugin.Execute(_downloadJob); + + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + DownloadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + + Assert.AreEqual(ur.download.full_path, "\\\\127.0.0.1\\C$\\Windows\\System32\\drivers\\etc\\hosts"); + } + + [TestMethod] + public void TestPathParsingUncWithFile() + { + string hostName = "127.0.0.1"; + string filePath = "C$\\Windows\\System32\\drivers\\etc"; + string fileName = "hosts"; + Dictionary downloadParams = new Dictionary() + { + { "path", filePath }, + { "file", fileName}, + { "host", hostName } + + }; + + _downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams); + _downloadPlugin.Execute(_downloadJob); + + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + DownloadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + + Assert.AreEqual(ur.download.full_path, "\\\\127.0.0.1\\C$\\Windows\\System32\\drivers\\etc\\hosts"); + } + string GetHashForFile(string filename) + { + using (var md5 = System.Security.Cryptography.MD5.Create()) + { + using (var stream = File.OpenRead(filename)) + { + var hash = md5.ComputeHash(stream); + var strHash = BitConverter.ToString(hash).Replace("-", "").ToLowerInvariant(); + Console.WriteLine(strHash); + return strHash; + } + + + } + } + string GetHashForByteArray(byte[] bytes) + { + using (var md5 = System.Security.Cryptography.MD5.Create()) + { + using (var stream = new MemoryStream(bytes)) + { + var hash = md5.ComputeHash(stream); + var strHash = BitConverter.ToString(hash).Replace("-", "").ToLowerInvariant(); + Console.WriteLine(strHash); + return strHash; + } + } + } } } diff --git a/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/UploadTests.cs b/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/UploadTests.cs index 35b41849..1a53705a 100644 --- a/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/UploadTests.cs +++ b/Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/UploadTests.cs @@ -39,11 +39,11 @@ public UploadTests() public void TestPathParsingLocalFull() { //Assert.IsTrue(false); - string fileName = Path.GetTempPath() + Guid.NewGuid().ToString() + ".txt"; - File.Create(fileName).Close(); + string fullPath = Path.GetTempPath() + Guid.NewGuid().ToString() + ".txt"; + File.Create(fullPath).Close(); Dictionary downloadParams = new Dictionary() { - {"path", fileName }, + {"path", fullPath }, }; _uploadJob.task.parameters = JsonSerializer.Serialize(downloadParams); @@ -53,17 +53,31 @@ public void TestPathParsingLocalFull() ((TestMessageManager)_messageManager).hasResponse.WaitOne(); UploadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); - File.Delete(fileName); + //Make sure - } - [TestMethod] - public void TestPathParsingUnc() - { - //Test to make sure the plugin parses local paths like we expect + Assert.AreEqual(ur.upload.full_path, fullPath); } [TestMethod] public void TestPathParsingRelative() { + //Assert.IsTrue(false); + string fileName = Guid.NewGuid().ToString() + ".txt"; + File.Create(fileName).Close(); + Dictionary downloadParams = new Dictionary() + { + {"path", ""}, + {"filename","myfile.txt" } + + }; + _uploadJob.task.parameters = JsonSerializer.Serialize(downloadParams); + _uploadPlugin.Execute(_uploadJob); + + ((TestMessageManager)_messageManager).hasResponse.WaitOne(); + UploadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); + + + //Make sure + Assert.AreEqual(ur.upload.full_path, Path.Combine(Directory.GetCurrentDirectory(),"myfile.txt")); //Test to make sure the plugin parses local paths like we expect } [TestMethod] @@ -192,10 +206,10 @@ public void TestFileNotExist() }; _uploadJob.task.parameters = JsonSerializer.Serialize(downloadParams); - //3 chunks - //Test to make sure the plugin parses local paths like we expect + _uploadPlugin.Execute(_uploadJob); + UploadResponse ur = JsonSerializer.Deserialize(((TestMessageManager)_messageManager).GetRecentOutput().Result); - //File.Delete(fileName); + Assert.IsTrue(ur.status == "error"); } public string TryHandleNextChunk(string path, int chunk) diff --git a/Payload_Type/athena/athena/agent_code/Agent/Agent.csproj b/Payload_Type/athena/athena/agent_code/Agent/Agent.csproj index 0e6f04a5..98b33e8a 100644 --- a/Payload_Type/athena/athena/agent_code/Agent/Agent.csproj +++ b/Payload_Type/athena/athena/agent_code/Agent/Agent.csproj @@ -35,11 +35,11 @@ - + CHECKYMANDERDEV - +