Merge pull request #27 from maxDcb/develop

Fix bug & small add

Author: maxDcb

C2Client/C2Client/AssistantPanel.py

@@ -135,33 +135,6 @@ def runCommand(self):
 
         else:
 
-            function_spec_assemblyExec = {
-                "name": "assemblyExec",
-                "description": "Execute a red team tool on a specific beacon using assembly execution",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "beacon_hash": {
-                            "type": "string",
-                            "description": "The unique hash identifying the beacon to execute the command on"
-                        },
-                        "listener_hash": {
-                            "type": "string",
-                            "description": "The unique hash identifying the listener at which the beacon is connected"
-                        },
-                        "tool": {
-                            "type": "string",
-                            "description": "The tool to use (e.g., Rubeus)"
-                        },
-                        "arguments": {
-                            "type": "string",
-                            "description": "Command line arguments to pass to the tool (e.g., '/dump')"
-                        }
-                    },
-                    "required": ["beacon_hash", "listener_hash", "tool"]
-                }
-            }
-
             function_spec_ls = {
                 "name": "ls",
                 "description": "List the contents of a specified directory on a specific beacon.",
@@ -250,7 +223,30 @@ def runCommand(self):
                 },
                     "required": ["beacon_hash", "listener_hash"]
             }
-
+            
+            function_spec_tree = {
+                "name": "tree",
+                "description": "Recursively display the directory structure of a specified path on a specific beacon in a tree-like format.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "beacon_hash": {
+                            "type": "string",
+                            "description": "The unique hash identifying the beacon to execute the command on"
+                        },
+                        "listener_hash": {
+                            "type": "string",
+                            "description": "The unique hash identifying the listener at which the beacon is connected"
+                        },
+                        "path": {
+                            "type": "string",
+                            "description": "The root directory path to start the tree traversal. If omitted, uses the current working directory.",
+                            "default": "."
+                        }
+                    },
+                    "required": ["beacon_hash", "listener_hash", "path"]
+                }
+            }
 
             api_key = os.environ.get("OPENAI_API_KEY")
 
@@ -275,7 +271,7 @@ def runCommand(self):
                         model="gpt-4o",
                         # model="gpt-3.5-turbo-1106", # test
                         messages=self.messages,
-                        functions=[function_spec_ls, function_spec_cd, function_spec_cat, function_spec_pwd],
+                        functions=[function_spec_ls, function_spec_cd, function_spec_cat, function_spec_pwd, function_spec_tree],
                         function_call="auto",
                         temperature=0.05
                     )
@@ -337,6 +333,16 @@ def executeCmd(self, cmd, args):
             if result.message:
                 self.printInTerminal("", commandLine, result.message.decode(encoding="latin1", errors="ignore"))
 
+        elif cmd == "tree":
+            beacon_hash = args["beacon_hash"]
+            listener_hash = args["listener_hash"]
+            path = args["path"]
+            commandLine = "tree " + path
+            command = TeamServerApi_pb2.Command(beaconHash=beacon_hash, listenerHash=listener_hash, cmd=commandLine)
+            result = self.grpcClient.sendCmdToSession(command)
+            if result.message:
+                self.printInTerminal("", commandLine, result.message.decode(encoding="latin1", errors="ignore"))
+
         elif cmd == "cd":
             beacon_hash = args["beacon_hash"]
             listener_hash = args["listener_hash"]

C2Client/C2Client/SessionPanel.py

@@ -115,7 +115,6 @@ def showContextMenu(self, position):
 
 
     # catch Interact and Stop menu click
-    # TODO add remove ?
     def  actionClicked(self, action):
         hash = self.item
         for ix, sessionStore in enumerate(self.listSessionObject):

C2Client/C2Client/TerminalPanel.py

@@ -108,6 +108,7 @@
 GrpcInfoListenerInstruction = "infoListener"
 GrpcBatcaveUploadToolInstruction = "batcaveUpload"
 GrpcSocksInstruction = "socks"
+GrpcReloadModulesInstruction = "reloadModules";
 
 BeaconFileWindows = "Beacon.exe"
 BeaconFileLinux = "Beacon"
@@ -154,13 +155,20 @@
 SetSubInstruction = "set"
 SearchSubInstruction = "search"
 
+ReloadModulesInstruction = "ReloadModules";
+ReloadModulesHelp = """ReloadModules:
+Command the TeamServer to reload the modules libraries located in TeamServerModulesDirectoryPath.
+Can be used to add a new functionality without restarting the TeamServer.
+"""
+
 
 def getHelpMsg():
     helpText  = HostInstruction+"\n"
     helpText += DropperInstruction+"\n"
     helpText += BatcaveInstruction+"\n"
     helpText += CredentialStoreInstruction+"\n"
-    helpText += SocksInstruction
+    helpText += SocksInstruction+"\n"
+    helpText += ReloadModulesInstruction
     return helpText
 
 completerData = [
@@ -176,7 +184,8 @@ def getHelpMsg():
             (GetSubInstruction, []),
             (SetSubInstruction, []),
             (SearchSubInstruction, [])
-             ])
+             ]),
+    (ReloadModulesInstruction,[]),
 ]
 
 InfoProcessing = "Processing..." 
@@ -275,6 +284,8 @@ def runCommand(self):
                         self.printInTerminal(commandLine, HostHelp)
                     elif instructions[1].lower() == CredentialStoreInstruction.lower():
                         self.printInTerminal(commandLine, CredentialStoreHelp)
+                    elif instructions[1].lower() == ReloadModulesInstruction.lower():
+                        self.printInTerminal(commandLine, ReloadModulesHelp)
                     elif instructions[1].lower() == DropperInstruction.lower():
                         availableModules = "- Available dropper:\n"
                         for module in DropperModules:
@@ -295,6 +306,8 @@ def runCommand(self):
                 self.runDropper(commandLine, instructions)
             elif instructions[0].lower()==SocksInstruction.lower():
                 self.runSocks(commandLine, instructions)
+            elif instructions[0].lower()==ReloadModulesInstruction.lower():
+                self.runReloadModules(commandLine, instructions)
             else:
                 self.printInTerminal(commandLine, ErrorCmdUnknow)
 
@@ -305,6 +318,16 @@ def runHelp(self):
         self.printInTerminal(HelpInstruction, getHelpMsg())
 
 
+    def runReloadModules(self, commandLine, instructions):
+        commandTeamServer = GrpcReloadModulesInstruction
+        termCommand = TeamServerApi_pb2.TermCommand(cmd=commandTeamServer)
+        resultTermCommand = self.grpcClient.sendTermCmd(termCommand)
+
+        result = resultTermCommand.result
+        self.printInTerminal(commandLine, result)
+        return   
+        
+
     def runSocks(self, commandLine, instructions):
         if len(instructions) < 2:
             self.printInTerminal(commandLine, SocksHelp)

core

@@ -1633,7 +1633,68 @@ grpc::Status TeamServer::SendTermCmd(grpc::ServerContext* context, const teamser
 	// TODO
 	else if(instruction==ReloadModulesInstruction)
 	{
-		// reload all the modules of the ../Modules directory
+		m_logger->info("Reloading TeamServer modules from directory: {0}", m_teamServerModulesDirectoryPath.c_str());
+
+		// Clear previously loaded modules
+			m_moduleCmd.clear();
+
+			try {
+				for (const auto& entry : fs::recursive_directory_iterator(m_teamServerModulesDirectoryPath)) 
+				{
+					if (fs::is_regular_file(entry.path()) && entry.path().extension() == ".so") 
+					{
+						m_logger->info("Trying to load {0}", entry.path().c_str());
+
+						void* handle = dlopen(entry.path().c_str(), RTLD_LAZY);
+						if (!handle) 
+						{
+							m_logger->warn("Failed to load {0}: {1}", entry.path().c_str(), dlerror());
+							continue;
+						}
+
+						// Derive constructor function name
+						std::string funcName = entry.path().filename();
+						funcName = funcName.substr(3); 							// remove lib
+						funcName = funcName.substr(0, funcName.length() - 3);	// remove .so
+						funcName += "Constructor";								// add Constructor
+
+						m_logger->info("Looking for constructor function: {0}", funcName);
+
+						constructProc construct = (constructProc)dlsym(handle, funcName.c_str());
+						if (!construct) {
+							m_logger->warn("Failed to find constructor: {0}", dlerror());
+							dlclose(handle);
+							continue;
+						}
+
+						ModuleCmd* moduleCmd = construct();
+						if (!moduleCmd) 
+						{
+							m_logger->warn("Constructor returned null");
+							dlclose(handle);
+							continue;
+						}
+
+						std::unique_ptr<ModuleCmd> moduleCmdPtr(moduleCmd);
+						moduleCmdPtr->setDirectories(
+							m_teamServerModulesDirectoryPath,
+							m_linuxModulesDirectoryPath,
+							m_windowsModulesDirectoryPath,
+							m_linuxBeaconsDirectoryPath,
+							m_windowsBeaconsDirectoryPath,
+							m_toolsDirectoryPath,
+							m_scriptsDirectoryPath
+						);
+
+						m_logger->info("Module {0} loaded", entry.path().filename().c_str());
+						m_moduleCmd.push_back(std::move(moduleCmdPtr));
+					}
+				}
+			} 
+			catch (const std::filesystem::filesystem_error& e) 
+			{
+				m_logger->warn("Error accessing module directory: {0}", e.what());
+			}
 	}
 	else if(instruction == SocksInstruction_)
 	{
@@ -1835,6 +1896,17 @@ int TeamServer::prepMsg(const std::string& input, C2Message& c2Message, bool isW
 				if (splitedCmd.size() == 2)
 				{
 					std::string param = splitedCmd[1];
+
+					// Handle the 4 historicals commands where the cmd name don't match the module file name
+					if(param=="ls")
+						param = "listDirectory";
+					else if(param=="cd")
+						param = "changeDirectory";
+					else if(param=="ps")
+						param = "listProcesses";
+					else if(param=="pwd")
+						param = "printWorkingDirectory";
+
 					if(param.size() >= 3 && param.substr(param.size() - 3) == ".so")
 					{