update unit tests

MythicAgents · Jan 14, 2025 · 5ac7656 · 5ac7656
1 parent a2bfe47
commit 5ac7656
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 44 deletions.
diff --git a/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/AgentTests/TaskingTests.cs b/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/AgentTests/TaskingTests.cs
@@ -80,7 +80,7 @@ public async Task TestGetTaskingMultiple()
             AthenaCore _agent = new AthenaCore(_profile, _taskManager, _logger, _config, _tokenManager, new List<IAgentMod>() { _agentMod });
             TestProfile prof = (TestProfile)_profile.First();
 
-            Task.Run(_agent.Start);
+            await _agent.Start();
             _profile.First().SetTaskingReceived += (sender, args) => taskingReceived.Set();
 
             ((TestTaskManager)_taskManager).WaitForNumberOfJobs(3);
@@ -109,18 +109,5 @@ public async Task TestGetTaskingNoTasks() {
             _profile.First().StopBeacon();
             Assert.IsTrue(((TestTaskManager)_taskManager).jobs.Count == 0);
         }
-        [TestMethod]
-        public async Task TestGetTaskingNullTasks()
-        {
-            ManualResetEvent taskingReceived = new ManualResetEvent(false);
-            IEnumerable<IProfile> _profile = new List<IProfile>() { new TestProfile(true) };
-            _profile.First().SetTaskingReceived += (sender, args) => taskingReceived.Set();
-            TestProfile prof = (TestProfile)_profile.First();
-            AthenaCore _agent = new AthenaCore(_profile, _taskManager, _logger, _config, _tokenManager, new List<IAgentMod>() { _agentMod });
-            Task.Run(_agent.Start);
-            prof.taskingSent.WaitOne(1000);
-            _profile.First().StopBeacon();
-            Assert.IsTrue(((TestTaskManager)_taskManager).jobs.Count == 0);
-        }
     }
 }
diff --git a/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/TestClasses/TestProfile.cs b/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/TestClasses/TestProfile.cs
@@ -17,38 +17,11 @@ public class TestProfile : IProfile
         private bool returnedTask = false;
         public CancellationTokenSource cts = new CancellationTokenSource();
         public ManualResetEvent taskingSent = new ManualResetEvent(false);
-        public TestProfile() {
-
-            checkinResponse = new CheckinResponse()
-            {
-                status = "success",
-                action = "checkin",
-                id = Guid.NewGuid().ToString(),
-                encryption_key = "",
-                decryption_key = "",
-                process_name = "",
-            };
-
-            getTaskingResponse = new GetTaskingResponse()
-            {
-                action = "get_tasking",
-                tasks = new List<ServerTask>()
-                {
-                    new ServerTask()
-                    {
-                        id = Guid.NewGuid().ToString(),
-                        command = "whoami",
-                        parameters = "",
-                    }
-                },
-                socks = new List<ServerDatagram>(),
-                rpfwd = new List<ServerDatagram>(),
-                delegates = new List<DelegateMessage>(),
-            };
+        public TestProfile() 
+        {
         }
         public TestProfile(CheckinResponse? checkinResponse)
         {
-            Console.WriteLine(checkinResponse.status);
             this.checkinResponse = checkinResponse;
         }
 
@@ -64,6 +37,18 @@ public TestProfile(bool nullable)
 
         public async Task<CheckinResponse> Checkin(Checkin checkin)
         {
+            if (this.checkinResponse is null)
+            {
+                return new CheckinResponse()
+                {
+                    status = "success",
+                    action = "checkin",
+                    id = Guid.NewGuid().ToString(),
+                    encryption_key = "",
+                    decryption_key = "",
+                    process_name = "",
+                };
+            }
             return this.checkinResponse;
         }
 
@@ -74,10 +59,36 @@ public async Task StartBeacon()
             {
                 if (!returnedTask)
                 {
-                    TaskingReceivedArgs tra = new TaskingReceivedArgs(getTaskingResponse);
+                    TaskingReceivedArgs tra;
+                    if (this.getTaskingResponse is null)
+                    {
+
+                        tra = new TaskingReceivedArgs(new GetTaskingResponse()
+                        {
+                            action = "get_tasking",
+                            tasks = new List<ServerTask>()
+                                {
+                                    new ServerTask()
+                                    {
+                                        id = Guid.NewGuid().ToString(),
+                                        command = "whoami",
+                                        parameters = "",
+                                    },
+                                },
+                            socks = new List<ServerDatagram>(),
+                            rpfwd = new List<ServerDatagram>(),
+                            delegates = new List<DelegateMessage>(),
+                        });
+                    }
+                    else
+                    {
+                        tra = new TaskingReceivedArgs(this.getTaskingResponse);
+                    }
+
                     SetTaskingReceived?.Invoke(this, tra);
                     returnedTask = true;
                     taskingSent.Set();
+                    return;
                 }
             }
             isRunning = false;