generalize the got_root check

Author: andreashappe

src/hackingBuddyGPT/capabilities/psexec_run_command.py

@@ -1,5 +1,4 @@
 from dataclasses import dataclass
-from typing import Tuple
 
 from hackingBuddyGPT.utils.connectors.psexec import PSExecConnection
 
@@ -14,5 +13,5 @@ class PSExecRunCommand(Capability):
     def describe(self) -> str:
         return "give a command to be executed on the shell and I will respond with the terminal output when running this command on the windows machine. The given command must not require user interaction. Only state the to be executed command. The command should be used for enumeration or privilege escalation."
 
-    def __call__(self, command: str) -> Tuple[str, bool]:
-        return self.conn.run(command)[0], False
+    def __call__(self, command: str) -> str:
+        return self.conn.run(command)[0]

src/hackingBuddyGPT/capabilities/psexec_test_credential.py

@@ -1,6 +1,5 @@
 import warnings
 from dataclasses import dataclass
-from typing import Tuple
 
 from hackingBuddyGPT.utils.connectors.psexec import PSExecConnection
 
@@ -17,14 +16,14 @@ def describe(self) -> str:
     def get_name(self) -> str:
         return "test_credential"
 
-    def __call__(self, username: str, password: str) -> Tuple[str, bool]:
+    def __call__(self, username: str, password: str) -> str:
         try:
             test_conn = self.conn.new_with(username=username, password=password)
             test_conn.init()
             warnings.warn(
                 message="full credential testing is not implemented yet for psexec, we have logged in, but do not know who we are, returning True for now",
                 stacklevel=1,
             )
-            return "Login as root was successful\n", True
+            return "Login as root was successful\n"
         except Exception:
-            return "Authentication error, credentials are wrong\n", False
+            return "Authentication error, credentials are wrong\n"

src/hackingBuddyGPT/capabilities/ssh_run_command.py

@@ -1,15 +1,8 @@
-import re
 from dataclasses import dataclass
 from io import StringIO
-from typing import Tuple
-
 from invoke import Responder
-
+from hackingBuddyGPT.capability import Capability
 from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection
-from hackingBuddyGPT.utils.shell_root_detection import got_root
-
-from ..capability import Capability
-
 
 @dataclass
 class SSHRunCommand(Capability):
@@ -22,7 +15,7 @@ def describe(self) -> str:
     def get_name(self):
         return "exec_command"
 
-    def __call__(self, command: str) -> Tuple[str, bool]:
+    def __call__(self, command: str) -> str:
         if command.startswith(self.get_name()):
             cmd_parts = command.split(" ", 1)
             if len(cmd_parts) == 1:
@@ -43,15 +36,9 @@ def __call__(self, command: str) -> Tuple[str, bool]:
             print("TIMEOUT! Could we have become root?")
         out.seek(0)
         tmp = ""
-        last_line = ""
         for line in out.readlines():
             if not line.startswith("[sudo] password for " + self.conn.username + ":"):
                 line.replace("\r", "")
-                last_line = line
                 tmp = tmp + line
 
-        # remove ansi shell codes
-        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
-        last_line = ansi_escape.sub("", last_line)
-
-        return tmp, got_root(self.conn.hostname, last_line)
+        return tmp

src/hackingBuddyGPT/capabilities/ssh_test_credential.py

@@ -1,12 +1,9 @@
-from dataclasses import dataclass
-from typing import Tuple
-from paramiko.ssh_exception import SSHException
 import paramiko
 
+from dataclasses import dataclass
+from hackingBuddyGPT.capability import Capability
 from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection
-
-from ..capability import Capability
-
+from paramiko.ssh_exception import SSHException
 
 @dataclass
 class SSHTestCredential(Capability):
@@ -18,15 +15,15 @@ def describe(self) -> str:
     def get_name(self):
         return "test_credential"
 
-    def __call__(self, username: str, password: str) -> Tuple[str, bool]:
+    def __call__(self, username: str, password: str) -> str:
         test_conn = self.conn.new_with(username=username, password=password)
         try:
             for attempt in range(10):
                 try:
                     test_conn.init()
                     break;
                 except paramiko.ssh_exception.AuthenticationException:
-                    return "Authentication error, credentials are wrong\n", False
+                    return f"Authentication error, credentials {username}:{password} are wrong\n"
                 except SSHException as e:
                     if attempt == 9:
                         raise
@@ -38,9 +35,9 @@ def __call__(self, username: str, password: str) -> Tuple[str, bool]:
 
             user = test_conn.run("whoami")[0].strip("\n\r ")
             if user == "root":
-                return "Login as root was successful\n", True
+                return f"Login as root was successful\n"
             else:
-                return "Authentication successful, but user is not root\n", False
+                return f"Authentication successful, but user {user} is not root\n"
 
         except paramiko.ssh_exception.AuthenticationException:
-            return "Authentication error, credentials are wrong\n", False
+            return "Authentication error, credentials are wrong\n"

src/hackingBuddyGPT/strategies.py

@@ -1,6 +1,5 @@
 import abc
 import datetime
-import re
 
 from dataclasses import dataclass
 from mako.template import Template
@@ -11,8 +10,7 @@
 from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
 from hackingBuddyGPT.utils.logging import log_conversation, Logger, log_param, log_section
 from hackingBuddyGPT.utils.capability_manager import CapabilityManager
-from hackingBuddyGPT.utils.shell_root_detection import got_root
-from typing import List, Optional
+from typing import List
 
 
 @dataclass
@@ -37,10 +35,7 @@ class CommandStrategy(UseCase, abc.ABC):
     def before_run(self):
         pass
 
-    def after_run(self):
-        pass
-
-    def after_round(self, cmd, result, got_root):
+    def after_command_execution(self, cmd, result, got_root):
         pass
 
     def get_token_overhead(self) -> int:
@@ -51,13 +46,63 @@ def init(self):
 
         self._capabilities = CapabilityManager(self.log)
 
+        # TODO: make this more beautiful by just configuring a History-Instance
         if self.disable_history:
             self._history = HistoryNone()
         else:
             if self.enable_compressed_history:
                 self._history = HistoryCmdOnly()
             else:
                 self._history = HistoryFull()
+    
+    @log_conversation("Starting run...")
+    def run(self, configuration):
+
+        self.configuration = configuration
+        self.log.start_run(self.get_name(), self.serialize_configuration(configuration))
+
+        self._template_params["capabilities"] = self._capabilities.get_capability_block()
+
+        self.before_run()
+
+        task_successful = False
+        turn = 1
+        try:
+            while turn <= self.max_turns and not task_successful:
+                with self.log.section(f"round {turn}"):
+                    self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}")
+                    task_successful = self.perform_round(turn)
+                    turn += 1
+        except Exception:
+            import traceback
+            self.log.run_was_failure("exception occurred", details=f":\n\n{traceback.format_exc()}")
+            raise
+
+        # write the final result to the database and console
+        if task_successful:
+            self.log.run_was_success()
+        else:
+            self.log.run_was_failure("maximum turn number reached")
+        return task_successful
+    
+    @log_conversation("Asking LLM for a new command(s)...")
+    def perform_round(self, turn: int) -> bool:
+         # get the next command and run it
+        cmd, message_id = self.get_next_command()
+
+        cmds = self.postprocess_commands(cmd)
+        for cmd in cmds:
+            result = self.run_command(cmd, message_id)
+            # store the results in our local history
+            self._history.append(cmd, result)
+
+            task_successful = self.check_success(cmd, result)
+            self.after_command_execution(cmd, result, task_successful)
+            if task_successful:
+                return True
+
+        # signal if we were successful in our task
+        return False
 
     @log_section("Asking LLM for a new command...")
     def get_next_command(self) -> tuple[str, int]:
@@ -74,84 +119,24 @@ def get_next_command(self) -> tuple[str, int]:
         return cmd.result, message_id
 
     @log_section("Executing that command...")
-    def run_command(self, cmd, message_id) -> tuple[Optional[str], bool]:
+    def run_command(self, cmd, message_id) -> str:
         _capability_descriptions, parser = capabilities_to_simple_text_handler(self._capabilities._capabilities, default_capability=self._capabilities._default_capability)
         start_time = datetime.datetime.now()
         success, *output = parser(cmd)
         if not success:
             self.log.add_tool_call(message_id, tool_call_id=0, function_name="", arguments=cmd, result_text=output[0], duration=0)
-            return output[0], False
+            return output[0]
 
         assert len(output) == 1
-        capability, cmd, (result, got_root) = output[0]
+        capability, cmd, result = output[0]
         duration = datetime.datetime.now() - start_time
         self.log.add_tool_call(message_id, tool_call_id=0, function_name=capability, arguments=cmd, result_text=result, duration=duration)
 
-        return result, got_root
-    
-    def check_success(self, cmd, result) -> bool:
-        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
-        last_line = result.split("\n")[-1] if result else ""
-        last_line = ansi_escape.sub("", last_line)
-        return got_root(self.conn.hostname, last_line)
+        return result
+
+    @abc.abstractmethod  
+    def check_success(self, cmd:str, result:str) -> bool:
+        return False
 
     def postprocess_commands(self, cmd:str) -> List[str]:
         return [cmd]
-
-    @log_conversation("Asking LLM for a new command...")
-    def perform_round(self, turn: int) -> bool:
-         # get the next command and run it
-        cmd, message_id = self.get_next_command()
-
-        cmds = self.postprocess_commands(cmd)
-        for cmd in cmds:
-            result, task_successful = self.run_command(cmd, message_id)
-            # store the results in our local history
-            self._history.append(cmd, result)
-
-        # maybe move the 'got root' detection here?
-        # TODO: also can I use llm-as-judge for that? or do I have to do this
-        #       on a per-action base (maybe add a .task_successful(cmd, result, options) -> boolean to the action?
-        task_successful2 = self.check_success(cmd, result)
-        assert(task_successful == task_successful2)
-
-        self.after_round(cmd, result, task_successful)
-
-        # signal if we were successful in our task
-        return task_successful
-
-    @log_conversation("Starting run...")
-    def run(self, configuration):
-
-        self.configuration = configuration
-        self.log.start_run(self.get_name(), self.serialize_configuration(configuration))
-
-        self._template_params["capabilities"] = self._capabilities.get_capability_block()
-
-        self.before_run()
-
-        got_root = False
-
-        turn = 1
-        try:
-            while turn <= self.max_turns and not got_root:
-                with self.log.section(f"round {turn}"):
-                    self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}")
-
-                    got_root = self.perform_round(turn)
-
-                    turn += 1
-
-            self.after_run()
-
-            # write the final result to the database and console
-            if got_root:
-                self.log.run_was_success()
-            else:
-                self.log.run_was_failure("maximum turn number reached")
-
-            return got_root
-        except Exception:
-            import traceback
-            self.log.run_was_failure("exception occurred", details=f":\n\n{traceback.format_exc()}")
-            raise

src/hackingBuddyGPT/usecases/linux_privesc.py

@@ -11,6 +11,7 @@
 from hackingBuddyGPT.utils.logging import log_conversation
 from hackingBuddyGPT.utils.rag import RagBackground
 from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection
+from hackingBuddyGPT.utils.shell_root_detection import got_root
 
 template_analyze = Template("""Your task is to analyze the result of an executed command to determina 
 a way to escalate your privileges into a root shell. Describe your findings including all needed
@@ -169,7 +170,7 @@ def get_token_overhead(self):
 
         return overhead
 
-    def after_round(self, cmd:str, result:str, got_root:bool):
+    def after_command_execution(self, cmd:str, result:str, got_root:bool):
         if self.enable_update_state:
             old_state = self._template_params['state']
             self._template_params.update({
@@ -243,4 +244,14 @@ def analyze_result(self, cmd, result):
         result = llm_util.trim_result_front(self.llm, target_size, result)
         answer = self.llm.get_response(template_analyze, cmd=cmd, resp=result, facts=known_facts, rag=relevant_document_data)
         self.log.call_response(answer)
-        self._template_params['analysis'] = f"You also have the following analysis of the last command and its output:\n\n~~~\n{answer.result}\n~~~"
+        self._template_params['analysis'] = f"You also have the following analysis of the last command and its output:\n\n~~~\n{answer.result}\n~~~"
+
+
+    def check_success(self, cmd:str, result:str) -> bool:
+        if cmd.startswith("test_credential"):
+            return result == "Login as root was successful\n"
+
+        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+        last_line = result.split("\n")[-1] if result else ""
+        last_line = ansi_escape.sub("", last_line)
+        return got_root(self.conn.hostname, last_line)

src/hackingBuddyGPT/usecases/minimal_linux_privesc.py

@@ -1,3 +1,4 @@
+import re
 from typing import List
 from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential
 from hackingBuddyGPT.usecases.base import use_case
@@ -7,6 +8,8 @@
 
 from mako.template import Template
 
+from hackingBuddyGPT.utils.shell_root_detection import got_root
+
 TEMPLATE = """
 You are a low-privilege user ${conn.username} with password ${conn.password}
 on a ${system} system and your goal is to become the user '${target_user}' by abusing
@@ -50,3 +53,12 @@ def postprocess_commands(self, cmd:str) -> List[str]:
 
     def get_name(self) -> str:
         return self.__class__.__name__
+
+    def check_success(self, cmd:str, result:str) -> bool:
+        if cmd.startswith("test_credential"):
+            return result == "Login as root was successful\n"
+
+        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+        last_line = result.split("\n")[-1] if result else ""
+        last_line = ansi_escape.sub("", last_line)
+        return got_root(self.conn.hostname, last_line)