mirror of
https://github.com/maxDcb/C2TeamServer
synced 2026-06-06 16:14:27 +00:00
Refine assistant tool orchestration
This commit is contained in:
Maxime dcb
+356
-243
@@ -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 = '<p style="white-space:pre">'+'<span style="color:blue;">['+now.strftime("%Y:%m:%d %H:%M:%S").rstrip()+']</span>'+'<span style="color:red;"> [+] </span>'+'<span style="color:red;">{}</span>'+'</p>'
|
||||
|
||||
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 = []
|
||||
|
||||
Reference in New Issue
Block a user