update assembly loading for tests

MythicAgents · Dec 28, 2024 · a2bfe47 · a2bfe47
1 parent 99cc6e4
commit a2bfe47
Showing 1 changed file with 19 additions and 32 deletions.
diff --git a/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/PluginLoader.cs b/Payload_Type/athena/athena/agent_code/Tests/Agent.Tests/PluginLoader.cs
@@ -9,6 +9,7 @@
 using Agent.Interfaces;
 using Agent.Models;
 using System.ComponentModel;
+using static IronPython.Modules._ast;
 
 namespace Agent.Tests
 {
@@ -22,46 +23,32 @@ public class PluginLoader
         public IPythonManager pyManager { get; set; } = new PythonManager();
         public IPlugin? LoadPluginFromDisk(string pluginName)
         {
-            string path = FindMostRecentPlugin(Path.Combine(Directory.GetCurrentDirectory(), "..", "..", "..", "..", ".."), pluginName);
-            //var path = Path.Combine(Directory.GetCurrentDirectory(), "..", "..", "..", "..", "..", pluginName, "bin", "LocalDebugHttp", "net8.0", $"{pluginName}.dll");
-            //if (!File.Exists(path))
-            //{
-            //    path = Path.Combine(Directory.GetCurrentDirectory(), "..", "..", "..", "..", "..", pluginName, "bin", "Debug", "net8.0", $"{pluginName}.dll");
-            //}
-            byte[] buf = File.ReadAllBytes(path);
-            Assembly asm = Assembly.Load(buf);
-
-            return ParseAssemblyForPlugin(asm, this.messageManager, this.agentConfig, this.logger, this.tokenManager, this.spawner, this.pyManager);
+            return GetPlugin(pluginName);
         }
-        private string FindMostRecentPlugin(string rootFolder, string pluginName)
+
+        private IPlugin? GetPlugin(string pluginName)
         {
-            if (string.IsNullOrEmpty(rootFolder) || !Directory.Exists(rootFolder))
+            try
+            {
+                Assembly _tasksAsm = Assembly.Load($"{pluginName}, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null");
+                foreach (Type t in _tasksAsm.GetTypes())
+                {
+                    if (typeof(IPlugin).IsAssignableFrom(t))
+                    {
+                        IPlugin plug = (IPlugin)Activator.CreateInstance(t, messageManager, agentConfig, logger, tokenManager, spawner, pyManager);
+                        return plug;
+                    }
+                }
+                return null;
+            }
+            catch
             {
-                throw new ArgumentException("The specified folder does not exist.");
+                return null;
             }
-
-            var latestDll = Directory.EnumerateFiles(rootFolder, $"{pluginName}.dll", SearchOption.AllDirectories)
-                                     .OrderByDescending(File.GetLastWriteTime)
-                                     .FirstOrDefault();
-
-            return latestDll;
         }
         public PluginLoader(IMessageManager messageManager)
         {
             this.messageManager = messageManager;
         }
-
-        private static IPlugin ParseAssemblyForPlugin(Assembly asm, IMessageManager messageManager, IAgentConfig agentConfig, ILogger logger, ITokenManager tokenManager, ISpawner spawner, IPythonManager pyManager)
-        {
-            foreach (Type t in asm.GetTypes())
-            {
-                if (typeof(IPlugin).IsAssignableFrom(t))
-                {
-                    IPlugin plug = (IPlugin)Activator.CreateInstance(t, messageManager, agentConfig, logger, tokenManager, spawner, pyManager);
-                    return plug;
-                }
-            }
-            return null;
-        }
     }
 }