From c8cab9c07634f38b048d975d05a35e2da202b9e9 Mon Sep 17 00:00:00 2001 From: Trinity-SYT-SECURITY Date: Tue, 6 Aug 2024 03:13:14 +0100 Subject: [PATCH 1/4] fix missing 1 required positional argument bug --- src/hackingBuddyGPT/usecases/agents.py | 49 ++++++++++++++++++-------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/src/hackingBuddyGPT/usecases/agents.py b/src/hackingBuddyGPT/usecases/agents.py index 003d4552..86807b87 100644 --- a/src/hackingBuddyGPT/usecases/agents.py +++ b/src/hackingBuddyGPT/usecases/agents.py @@ -55,9 +55,8 @@ def set_initial_state(self, initial_state:AgentWorldview): def set_template(self, template:str): self._template = Template(filename=template) self._template_size = self.llm.count_tokens(self._template.source) - - def perform_round(self, turn:int) -> bool: - got_root : bool = False + def perform_round(self, turn: int) -> bool:##fix code + got_root: bool = False with self.console.status("[bold green]Asking LLM for a new command..."): # TODO output/log state @@ -71,15 +70,35 @@ def perform_round(self, turn:int) -> bool: cmd = llm_util.cmd_output_fixer(answer.result) with self.console.status("[bold green]Executing that command..."): - self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) - capability = self.get_capability(cmd.split(" ", 1)[0]) - result, got_root = capability(cmd) - - # log and output the command and its result - self.log_db.add_log_query(self._run_id, turn, cmd, result, answer) - self._state.update(capability, cmd, result) - # TODO output/log new state - self.console.print(Panel(result, title=f"[bold cyan]{cmd}")) - - # if we got root, we can stop the loop - return got_root + self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) + + # Assuming command is of the form "capability_name arg1 arg2" + parts = cmd.split(" ", 1) + if len(parts) == 2: + capability_name, args = parts + capability = self.get_capability(capability_name) + + if capability: + # Assuming capability requires multiple arguments + # Adjust the argument unpacking based on capability's requirements + args_list = args.split() # Split arguments into a list + + try: + result, got_root = capability(*args_list) + except TypeError as e: + result = f"Error executing command: {e}" + got_root = False + else: + result = f"Unknown capability: {capability_name}" + got_root = False + else: + result = "Command format error. Expected 'capability_name arg1 arg2'." + got_root = False + + # Log and output the command and its result + self.log_db.add_log_query(self._run_id, turn, cmd, result, answer) + self._state.update(capability, cmd, result) # Assuming capability is available + self.console.print(Panel(result, title=f"[bold cyan]{cmd}")) + + # If we got root, we can stop the loop + return got_root From 6644f996e0802ea1ddeed3e07f9462f7f40d3f56 Mon Sep 17 00:00:00 2001 From: Trinity-SYT-SECURITY Date: Tue, 6 Aug 2024 03:20:45 +0100 Subject: [PATCH 2/4] fix missing 1 required positional argument bug --- src/hackingBuddyGPT/usecases/minimal/agent.py | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/src/hackingBuddyGPT/usecases/minimal/agent.py b/src/hackingBuddyGPT/usecases/minimal/agent.py index 555a0684..fe35477a 100644 --- a/src/hackingBuddyGPT/usecases/minimal/agent.py +++ b/src/hackingBuddyGPT/usecases/minimal/agent.py @@ -26,7 +26,6 @@ def init(self): self.add_capability(SSHRunCommand(conn=self.conn), default=True) self.add_capability(SSHTestCredential(conn=self.conn)) self._template_size = self.llm.count_tokens(template_next_cmd.source) - def perform_round(self, turn): got_root : bool = False @@ -39,13 +38,27 @@ def perform_round(self, turn): cmd = llm_util.cmd_output_fixer(answer.result) with self.console.status("[bold green]Executing that command..."): - self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) - result, got_root = self.get_capability(cmd.split(" ", 1)[0])(cmd) + self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:")) + + # Assuming cmd is of the form "username password" + parts = cmd.split(" ", 1) + if len(parts) == 2: + username, password = parts + ##here fix! + result, got_root = self.get_capability("test_credential")(username, password) + else: + # Handle other cases or log error + result = "Command format error. Expected 'username password'." + got_root = False - # log and output the command and its result - self.log_db.add_log_query(self._run_id, turn, cmd, result, answer) - self._sliding_history.add_command(cmd, result) - self.console.print(Panel(result, title=f"[bold cyan]{cmd}")) + #self.log_db.add_log_query(self._run_id, cmd, result, answer) + self.log_db.add_log_query(self._run_id, turn, cmd, result, answer) + self._sliding_history.add_command(cmd, result) + self.console.print(Panel(result, title=f"[bold cyan]{cmd}")) - # if we got root, we can stop the loop return got_root + + + + + From 70ea9836a48f2b1e3aecc061a7d1cf05e169cc82 Mon Sep 17 00:00:00 2001 From: Trinity-SYT-SECURITY Date: Tue, 6 Aug 2024 03:25:52 +0100 Subject: [PATCH 3/4] fix sqlite3.operationalerror: database is locked issue --- .../utils/db_storage/db_storage.py | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/src/hackingBuddyGPT/utils/db_storage/db_storage.py b/src/hackingBuddyGPT/utils/db_storage/db_storage.py index 06787e4e..fd1523ea 100644 --- a/src/hackingBuddyGPT/utils/db_storage/db_storage.py +++ b/src/hackingBuddyGPT/utils/db_storage/db_storage.py @@ -11,10 +11,19 @@ def __init__(self, connection_string: str = parameter(desc="sqlite3 database con def init(self): self.connect() self.setup_db() - + + def connect(self): +# self.db = sqlite3.connect(self.connection_string, timeout=10) # Set timeout to 10 seconds + self.db = sqlite3.connect(self.connection_string, check_same_thread=False, timeout=10) + self.cursor = self.db.cursor() + + + ''' def connect(self): - self.db = sqlite3.connect(self.connection_string) + self.db = sqlite3.connect(self.connection_string, timeout=10.0) + self.cursor = self.db.cursor() + ''' def insert_or_select_cmd(self, name: str) -> int: results = self.cursor.execute("SELECT id, name FROM commands WHERE name = ?", (name,)).fetchall() @@ -80,7 +89,7 @@ def setup_db(self): self.query_cmd_id = self.insert_or_select_cmd('query_cmd') self.analyze_response_id = self.insert_or_select_cmd('analyze_response') self.state_update_id = self.insert_or_select_cmd('update_state') - + ''' def create_new_run(self, model, context_size, tag): self.cursor.execute( "INSERT INTO runs (model, context_size, state, tag, started_at) VALUES (?, ?, ?, ?, datetime('now'))", @@ -93,7 +102,21 @@ def add_log_query(self, run_id, round, cmd, result, answer): ( run_id, round, self.query_cmd_id, cmd, result, answer.duration, answer.tokens_query, answer.tokens_response, answer.prompt, answer.answer)) + ''' + def create_new_run(self, model, context_size, tag): + with self.db: + self.cursor.execute( + "INSERT INTO runs (model, context_size, state, tag, started_at) VALUES (?, ?, ?, ?, datetime('now'))", + (model, context_size, "in progress", tag)) + return self.cursor.lastrowid + + def add_log_query(self, run_id, round, cmd, result, answer): + with self.db: + self.cursor.execute( + "INSERT INTO queries (run_id, round, cmd_id, query, response, duration, tokens_query, tokens_response, prompt, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + (run_id, round, self.query_cmd_id, cmd, result, answer.duration, answer.tokens_query, answer.tokens_response, answer.prompt, answer.answer)) + def add_log_analyze_response(self, run_id, round, cmd, result, answer): self.cursor.execute( "INSERT INTO queries (run_id, round, cmd_id, query, response, duration, tokens_query, tokens_response, prompt, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", @@ -206,4 +229,4 @@ def run_was_failure(self, run_id, round): self.db.commit() def commit(self): - self.db.commit() + self.db.commit() \ No newline at end of file From 66340e3db2d382e11101d474df5018ebb9a6c5ad Mon Sep 17 00:00:00 2001 From: Trinity-SYT-SECURITY Date: Tue, 6 Aug 2024 03:27:44 +0100 Subject: [PATCH 4/4] fix sqlite3.operationalerror: database is locked issue --- .../utils/db_storage/db_storage.py | 22 +------------------ 1 file changed, 1 insertion(+), 21 deletions(-) diff --git a/src/hackingBuddyGPT/utils/db_storage/db_storage.py b/src/hackingBuddyGPT/utils/db_storage/db_storage.py index fd1523ea..907b9af9 100644 --- a/src/hackingBuddyGPT/utils/db_storage/db_storage.py +++ b/src/hackingBuddyGPT/utils/db_storage/db_storage.py @@ -17,13 +17,6 @@ def connect(self): self.db = sqlite3.connect(self.connection_string, check_same_thread=False, timeout=10) self.cursor = self.db.cursor() - - ''' - def connect(self): - self.db = sqlite3.connect(self.connection_string, timeout=10.0) - - self.cursor = self.db.cursor() - ''' def insert_or_select_cmd(self, name: str) -> int: results = self.cursor.execute("SELECT id, name FROM commands WHERE name = ?", (name,)).fetchall() @@ -89,20 +82,7 @@ def setup_db(self): self.query_cmd_id = self.insert_or_select_cmd('query_cmd') self.analyze_response_id = self.insert_or_select_cmd('analyze_response') self.state_update_id = self.insert_or_select_cmd('update_state') - ''' - def create_new_run(self, model, context_size, tag): - self.cursor.execute( - "INSERT INTO runs (model, context_size, state, tag, started_at) VALUES (?, ?, ?, ?, datetime('now'))", - (model, context_size, "in progress", tag)) - return self.cursor.lastrowid - - def add_log_query(self, run_id, round, cmd, result, answer): - self.cursor.execute( - "INSERT INTO queries (run_id, round, cmd_id, query, response, duration, tokens_query, tokens_response, prompt, answer) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ( - run_id, round, self.query_cmd_id, cmd, result, answer.duration, answer.tokens_query, answer.tokens_response, - answer.prompt, answer.answer)) - ''' + def create_new_run(self, model, context_size, tag): with self.db: self.cursor.execute(