Skip to content

Commit

Permalink
added unit tests to Upload/download
Browse files Browse the repository at this point in the history
  • Loading branch information
checkymander committed Feb 7, 2024
1 parent e8807a9 commit 2c46c39
Show file tree
Hide file tree
Showing 3 changed files with 161 additions and 19 deletions.
138 changes: 133 additions & 5 deletions Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/DownloadTests.cs
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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<string, string> downloadParams = new Dictionary<string, string>()
{
{ "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<DownloadResponse>(((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<DownloadResponse>(((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<string, string> downloadParams = new Dictionary<string, string>()
{
{ "host", Dns.GetHostName() },
{ "path", fileName },

};
_downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams);
_downloadPlugin.Execute(_downloadJob);

((TestMessageManager)_messageManager).hasResponse.WaitOne();
DownloadResponse ur = JsonSerializer.Deserialize<DownloadResponse>(((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<DownloadResponse>(((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<string, string> downloadParams = new Dictionary<string, string>()
{
{"host", hostName },
{"path", filePath },

};
_downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams);
_downloadPlugin.Execute(_downloadJob);

((TestMessageManager)_messageManager).hasResponse.WaitOne();
DownloadResponse ur = JsonSerializer.Deserialize<DownloadResponse>(((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<string, string> downloadParams = new Dictionary<string, string>()
{
{ "path", filePath },
{ "file", fileName},
{ "host", hostName }

};

_downloadJob.task.parameters = JsonSerializer.Serialize(downloadParams);
_downloadPlugin.Execute(_downloadJob);

((TestMessageManager)_messageManager).hasResponse.WaitOne();
DownloadResponse ur = JsonSerializer.Deserialize<DownloadResponse>(((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;
}
}
}
}
}
38 changes: 26 additions & 12 deletions Payload_Type/athena/athena/agent_code/Agent.Tests/PluginTests/UploadTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, string> downloadParams = new Dictionary<string, string>()
{
{"path", fileName },
{"path", fullPath },

};
_uploadJob.task.parameters = JsonSerializer.Serialize(downloadParams);
Expand All @@ -53,17 +53,31 @@ public void TestPathParsingLocalFull()
((TestMessageManager)_messageManager).hasResponse.WaitOne();
UploadResponse ur = JsonSerializer.Deserialize<UploadResponse>(((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<string, string> downloadParams = new Dictionary<string, string>()
{
{"path", ""},
{"filename","myfile.txt" }

};
_uploadJob.task.parameters = JsonSerializer.Serialize(downloadParams);
_uploadPlugin.Execute(_uploadJob);

((TestMessageManager)_messageManager).hasResponse.WaitOne();
UploadResponse ur = JsonSerializer.Deserialize<UploadResponse>(((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]
Expand Down Expand Up @@ -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<UploadResponse>(((TestMessageManager)_messageManager).GetRecentOutput().Result);

//File.Delete(fileName);
Assert.IsTrue(ur.status == "error");
}

public string TryHandleNextChunk(string path, int chunk)
Expand Down
4 changes: 2 additions & 2 deletions Payload_Type/athena/athena/agent_code/Agent/Agent.csproj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,11 +35,11 @@



<PropertyGroup Condition="'$(COMPUTERNAME)' == 'DESKTOP-GRJNOH2' Or '$(COMPUTERNAME)' == 'DEV3' Or '$(COMPUTERNAME)' == 'DEV1'">
<PropertyGroup Condition="'$(COMPUTERNAME)' == 'DESKTOP-GRJNOH2' Or '$(COMPUTERNAME)' == 'DEV3' Or '$(COMPUTERNAME)' == 'DESKTOP-FIJBR89'">
<DefineConstants>CHECKYMANDERDEV</DefineConstants>
</PropertyGroup>

<ItemGroup Condition="'$(COMPUTERNAME)' == 'DESKTOP-GRJNOH2' Or '$(COMPUTERNAME)' == 'DEV3' Or '$(COMPUTERNAME)' == 'DEV1'">
<ItemGroup Condition="'$(COMPUTERNAME)' == 'DESKTOP-GRJNOH2' Or '$(COMPUTERNAME)' == 'DEV3' Or '$(COMPUTERNAME)' == 'DESKTOP-FIJBR89'">
<ProjectReference Include="..\Agent.Profiles.Http\Agent.Profiles.Http.csproj" />
<ProjectReference Include="..\Agent.Crypto.Aes\Agent.Crypto.Aes.csproj" />
<ProjectReference Include="..\Agent.Managers.Windows\Agent.Managers.Windows.csproj" />
Expand Down

0 comments on commit 2c46c39

Please sign in to comment.