updated download test for single chunk donwload

MythicAgents · Feb 7, 2024 · 6beaf02 · 6beaf02
1 parent 2c46c39
commit 6beaf02
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
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
@@ -96,12 +96,15 @@ public void TestSingleChunkDownload()
 
             ((TestMessageManager)_messageManager).hasResponse.WaitOne();
             DownloadResponse ur = JsonSerializer.Deserialize<DownloadResponse>(((TestMessageManager)_messageManager).GetRecentOutput().Result);
+
+            Assert.AreNotEqual(ur.status, "error");
             ServerResponseResult responseResult = new ServerResponseResult()
             {
                 task_id = "123",
                 file_id = "1234",
-                total_chunks = 4,
-                chunk_num = 1,
+                total_chunks = ur.download.total_chunks,
+                chunk_num = ur.download.chunk_num,
+
                 status = "success"
             };
             _downloadPlugin.HandleNextMessage(responseResult);
@@ -140,9 +143,9 @@ public void TestHandleNextChunkFailure()
 
             _downloadPlugin.HandleNextMessage(responseResult);
             ((TestMessageManager)_messageManager).hasResponse.WaitOne();
-            ur = JsonSerializer.Deserialize<DownloadResponse>(((TestMessageManager)_messageManager).GetRecentOutput().Result);
+            string output = ((TestMessageManager)_messageManager).GetRecentOutput().Result;
 
-            Assert.AreEqual(ur.status, "error");
+            Assert.AreEqual(output, "An error occurred while communicating with the server." + Environment.NewLine);
             //Test to make sure the plugin parses local paths like we expect
         }
 

diff --git a/Payload_Type/athena/athena/agent_code/Agent.Tests/TestClasses/TestMessageManager.cs b/Payload_Type/athena/athena/agent_code/Agent.Tests/TestClasses/TestMessageManager.cs
@@ -132,7 +132,7 @@ public async Task WriteLine(string? output, string task_id, bool completed)
 
         public async Task<string> GetRecentOutput()
         {
-            return taskResponses.FirstOrDefault();
+            return taskResponses.LastOrDefault();
         }
 
         public Task AddResponse(InteractMessage im)