diff --git a/C2Client/C2Client/AssistantPanel.py b/C2Client/C2Client/AssistantPanel.py index 7516096..80d8ebe 100644 --- a/C2Client/C2Client/AssistantPanel.py +++ b/C2Client/C2Client/AssistantPanel.py @@ -76,6 +76,15 @@ You also point out security gaps that could be leveraged. You understand operati # Maximum number of messages to retain self.MAX_MESSAGES = 20 + # Track pending tool execution state + self.awaiting_tool_result = False + self.pending_tool_name = None + self.pending_tool_context = None + self.tool_call_count = 0 + self.max_function_calls = 5 + + self._openai_client = None + def nextCompletion(self): index = self._compl.currentIndex() @@ -104,13 +113,67 @@ You also point out security gaps that could be leveraged. You understand operati def consoleAssistantMethod(self, action, beaconHash, listenerHash, context, cmd, result): - if action == "receive": - # print("consoleAssistantMethod", "-Context:\n" + context + "\n\n-Command sent:\n" + cmd + "\n\n-Response:\n" + result) - self.messages.append({"role": "user", "content": cmd + "\n" + result}) - elif action == "send": - toto = 1 + if action != "receive": + return + command_text = cmd or "" + if isinstance(command_text, bytes): + command_text = command_text.decode("latin1", errors="ignore") + command_text = command_text.strip() + output_text = result or "" + if isinstance(output_text, bytes): + output_text = output_text.decode("latin1", errors="ignore") + output_text = output_text.replace(chr(0), "") + + display_output = output_text if output_text else "[no output]" + + awaiting_result = False + if self.awaiting_tool_result: + if self.pending_tool_context: + awaiting_result = ( + self.pending_tool_context.get("beacon_hash") == beaconHash + and self.pending_tool_context.get("listener_hash") == listenerHash + ) + else: + awaiting_result = True + + if awaiting_result: + header = command_text or "[assistant command]" + self.printInTerminal("Command:", header, display_output) + + function_name = self.pending_tool_name or "unknown" + self.messages.append({"role": "function", "name": function_name, "content": display_output}) + self._trim_message_history() + + self.awaiting_tool_result = False + self.pending_tool_name = None + self.pending_tool_context = None + self.tool_call_count += 1 + + try: + self._request_assistant_response() + except openai.APIConnectionError as e: + print(f"Server connection error: {e.__cause__}") + except openai.RateLimitError as e: + print(f"OpenAI RATE LIMIT error {e.status_code}: {e.response}") + except openai.APIStatusError as e: + print(f"OpenAI STATUS error {e.status_code}: {e.response}") + except openai.BadRequestError as e: + print(f"OpenAI BAD REQUEST error {e.status_code}: {e.response}") + except Exception as e: + print(f"An unexpected error occurred: {e}") + else: + combined = command_text + if output_text: + combined = f"{command_text}\n{output_text}" if command_text else output_text + + if combined.strip(): + self.messages.append({"role": "user", "content": combined}) + self._trim_message_history() + + header = command_text or "[command]" + self.printInTerminal("Command:", header, display_output) def event(self, event): if event.type() == QEvent.Type.KeyPress and event.key() == Qt.Key.Key_Tab: self.tabPressed.emit() @@ -118,18 +181,21 @@ You also point out security gaps that could be leveraged. You understand operati return super().event(event) - def printInTerminal(self, cmd, result): + def printInTerminal(self, header="", message="", detail=""): now = datetime.now() formater = '

'+'['+now.strftime("%Y:%m:%d %H:%M:%S").rstrip()+']'+' [+] '+'{}'+'

' self.sem.acquire() - if cmd: - self.editorOutput.appendHtml(formater.format(cmd)) - self.editorOutput.insertPlainText("\n") - if result: - self.editorOutput.insertPlainText(result) - self.editorOutput.insertPlainText("\n") - self.sem.release() + try: + if header: + self.editorOutput.appendHtml(formater.format(header)) + self.editorOutput.insertPlainText("\n") + for text in (message, detail): + if text: + self.editorOutput.insertPlainText(text) + self.editorOutput.insertPlainText("\n") + finally: + self.sem.release() def runCommand(self): @@ -139,251 +205,291 @@ You also point out security gaps that could be leveraged. You understand operati if commandLine == "": self.printInTerminal("", "") + return - else: + if self.awaiting_tool_result: + self.printInTerminal("Analysis:", "Waiting for previous command output before continuing.") + return - function_spec_ls = { - "name": "ls", - "description": "List the contents of a specified directory on a specific beacon.", - "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 path of the directory to list. If omitted, uses the current working directory.", - "default": "." - } - }, - "required": ["beacon_hash", "listener_hash", "path"] - } - } + client = self._get_openai_client() + if client is None: + self.printInTerminal("OPENAI_API_KEY is not set, functionality deactivated.", "") + return - function_spec_cd = { - "name": "cd", - "description": "Change the working directory for subsequent module execution or file operations on a specific beacon.", - "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": "Absolute or relative path to change to (e.g., '../modules', '/tmp')." - } - }, - "required": ["beacon_hash", "listener_hash", "path"] - } - } + # Reset state for a new round of tool calls triggered by operator input + self.awaiting_tool_result = False + self.pending_tool_name = None + self.pending_tool_context = None + self.tool_call_count = 0 - function_spec_cat = { - "name": "cat", - "description": "Read and return the contents of a specified file on disk on a specific beacon.", - "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": "Absolute or relative path to the file (e.g., './modules/shellcode.bin', '/etc/hosts')." - } - }, - "required": ["beacon_hash", "listener_hash", "path"] - } - } + # Add user command to the conversation history + self.messages.append({"role": "user", "content": commandLine}) + self._trim_message_history() - function_spec_pwd = { - "name": "pwd", - "description": "Return the current working directory path on a specific beacon.", - "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" - } - } - }, - "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") - - if api_key: - client = OpenAI( - # This is the default and can be omitted - api_key=api_key, - ) - - # Add user command output - self.messages.append({"role": "user", "content": commandLine}) - - if len(self.messages) > self.MAX_MESSAGES * 2 + 1: - # Always keep the first message (system prompt) - system_prompt = self.messages[0] - recent_messages = self.messages[-(self.MAX_MESSAGES * 2):] - self.messages = [system_prompt] + recent_messages - - try: - # Call OpenAI API - response = client.chat.completions.create( - 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, function_spec_tree], - function_call="auto", - temperature=0.05 - ) - - self.printInTerminal("User:", commandLine) - - message = response.choices[0].message - # print(message) - - function_call = message.function_call - if function_call: - name = function_call.name - args = json.loads(function_call.arguments) - # print(f"Model wants to call `{name}` with arguments: {args}") - - self.printInTerminal("Analysis:", f"Model wants to call `{name}` with arguments: {args}") - - self.executeCmd(name, args) - - - assistant_reply = message.content - if assistant_reply: - self.printInTerminal("Analysis:", assistant_reply) - - # Add assistant's response to conversation - self.messages.append({"role": "assistant", "content": assistant_reply}) - - except openai.APIConnectionError as e: - print(f"Server connection error: {e.__cause__}") - - except openai.RateLimitError as e: - print(f"OpenAI RATE LIMIT error {e.status_code}: {e.response}") - - except openai.APIStatusError as e: - print(f"OpenAI STATUS error {e.status_code}: {e.response}") - - except openai.BadRequestError as e: - print(f"OpenAI BAD REQUEST error {e.status_code}: {e.response}") - - except Exception as e: - print(f"An unexpected error occurred: {e}") - - else: - self.printInTerminal("OPENAI_API_KEY is not set, functionality deactivated.", "") - + try: + self.printInTerminal("User:", commandLine) + self._request_assistant_response() + except openai.APIConnectionError as e: + print(f"Server connection error: {e.__cause__}") + except openai.RateLimitError as e: + print(f"OpenAI RATE LIMIT error {e.status_code}: {e.response}") + except openai.APIStatusError as e: + print(f"OpenAI STATUS error {e.status_code}: {e.response}") + except openai.BadRequestError as e: + print(f"OpenAI BAD REQUEST error {e.status_code}: {e.response}") + except Exception as e: + print(f"An unexpected error occurred: {e}") self.setCursorEditorAtEnd() + def _get_openai_client(self): + if self._openai_client is not None: + return self._openai_client + + api_key = os.environ.get("OPENAI_API_KEY") + if not api_key: + return None + + self._openai_client = OpenAI(api_key=api_key) + return self._openai_client + + def _function_specs(self): + function_spec_ls = { + "name": "ls", + "description": "List the contents of a specified directory on a specific beacon.", + "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 path of the directory to list. If omitted, uses the current working directory.", + "default": "." + } + }, + "required": ["beacon_hash", "listener_hash", "path"] + } + } + + function_spec_cd = { + "name": "cd", + "description": "Change the working directory for subsequent module execution or file operations on a specific beacon.", + "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": "Absolute or relative path to change to (e.g., '../modules', '/tmp').", + } + }, + "required": ["beacon_hash", "listener_hash", "path"] + } + } + + function_spec_cat = { + "name": "cat", + "description": "Read and return the contents of a specified file on disk on a specific beacon.", + "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": "Absolute or relative path to the file (e.g., './modules/shellcode.bin', '/etc/hosts')." + } + }, + "required": ["beacon_hash", "listener_hash", "path"] + } + } + + function_spec_pwd = { + "name": "pwd", + "description": "Return the current working directory path on a specific beacon.", + "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" + } + }, + "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"] + } + } + + return [ + function_spec_ls, + function_spec_cd, + function_spec_cat, + function_spec_pwd, + function_spec_tree, + ] + + def _request_assistant_response(self): + if self.awaiting_tool_result: + return + + client = self._get_openai_client() + if client is None: + return + + response = client.chat.completions.create( + model="gpt-4o", + messages=self.messages, + functions=self._function_specs(), + function_call="auto", + temperature=0.05, + ) + + message = response.choices[0].message + function_call = getattr(message, "function_call", None) + + if function_call and function_call.name: + if self.tool_call_count >= self.max_function_calls: + warning = "Maximum number of tool calls reached without final response." + self.printInTerminal("Analysis:", warning) + self.messages.append({"role": "user", "content": warning}) + self._trim_message_history() + return + + self._handle_function_call(message) + return + + assistant_reply = message.content + if assistant_reply: + self.printInTerminal("Analysis:", assistant_reply) + self.messages.append({"role": "assistant", "content": assistant_reply}) + self._trim_message_history() + + def _handle_function_call(self, message): + function_call = message.function_call + name = function_call.name + raw_arguments = function_call.arguments or "{}" + + self.messages.append({ + "role": message.role, + "content": message.content, + "function_call": { + "name": name, + "arguments": raw_arguments, + }, + }) + self._trim_message_history() + + self.pending_tool_context = None + + try: + args = json.loads(raw_arguments) + except json.JSONDecodeError as decode_error: + error_message = f"Error decoding arguments for `{name}`: {decode_error}" + self.printInTerminal("Analysis:", error_message) + self.messages.append({"role": "function", "name": name, "content": error_message}) + self._trim_message_history() + self.tool_call_count += 1 + return + + step_index = self.tool_call_count + 1 + step_info = f"Step {step_index}: calling `{name}` with arguments: {args}" + self.printInTerminal("Analysis:", step_info) + + try: + self.executeCmd(name, args) + except Exception as tool_error: + tool_error_message = f"Error executing `{name}`: {tool_error}" + self.printInTerminal("Analysis:", tool_error_message) + self.messages.append({"role": "function", "name": name, "content": tool_error_message}) + self._trim_message_history() + self.tool_call_count += 1 + return + + self.awaiting_tool_result = True + self.pending_tool_name = name + self.pending_tool_context = { + "beacon_hash": args.get("beacon_hash"), + "listener_hash": args.get("listener_hash"), + } def executeCmd(self, cmd, args): - - if cmd == "ls": - beacon_hash = args["beacon_hash"] - listener_hash = args["listener_hash"] - path = args["path"] - commandLine = "ls " + 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")) + supported_commands = {"ls", "tree", "cd", "cat", "pwd"} + if cmd not in supported_commands: + raise ValueError(f"Unsupported command type: {cmd}") - 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")) + required_keys = ["beacon_hash", "listener_hash"] + if cmd != "pwd": + required_keys.append("path") - elif cmd == "cd": - beacon_hash = args["beacon_hash"] - listener_hash = args["listener_hash"] - path = args["path"] - commandLine = "cd " + 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")) + missing = [key for key in required_keys if key not in args] + if missing: + raise KeyError(f"Missing required argument(s) for `{cmd}`: {', '.join(missing)}") - elif cmd == "cat": - beacon_hash = args["beacon_hash"] - listener_hash = args["listener_hash"] - path = args["path"] - commandLine = "cat " + 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 == "pwd": - beacon_hash = args["beacon_hash"] - listener_hash = args["listener_hash"] - commandLine = "pwd" - 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")) + beacon_hash = args["beacon_hash"] + listener_hash = args["listener_hash"] + if cmd == "pwd": + command_line = "pwd" else: - raise ValueError("Unsupported command type") - - return - + path = args["path"] + command_line = f"{cmd} {path}" + + command = TeamServerApi_pb2.Command( + beaconHash=beacon_hash, + listenerHash=listener_hash, + cmd=command_line, + ) + self.grpcClient.sendCmdToSession(command) + + return command_line + # setCursorEditorAtEnd def setCursorEditorAtEnd(self): @@ -392,6 +498,13 @@ You also point out security gaps that could be leveraged. You understand operati self.editorOutput.setTextCursor(cursor) + def _trim_message_history(self): + if len(self.messages) > self.MAX_MESSAGES * 2 + 1: + system_prompt = self.messages[0] + recent_messages = self.messages[-(self.MAX_MESSAGES * 2):] + self.messages = [system_prompt] + recent_messages + + class CommandEditor(QLineEdit): tabPressed = pyqtSignal() cmdHistory = []