Merge pull request #132 from ipa-lab/restructuring

Restructuring

Author: andreashappe

src/hackingBuddyGPT/strategies.py

@@ -1,7 +1,7 @@
 import abc
 from dataclasses import dataclass
 import datetime
-from typing import Optional
+from typing import List, Optional
 import re
 
 from mako.template import Template
@@ -65,7 +65,7 @@ def get_next_command(self) -> tuple[str, int]:
         cmd = self.llm.get_response(self._template, **self._template_params)
         message_id = self.log.call_response(cmd)
 
-        return llm_util.cmd_output_fixer(cmd.result), message_id
+        return cmd.result, message_id
 
     @log_section("Executing that command...")
     def run_command(self, cmd, message_id) -> tuple[Optional[str], bool]:
@@ -89,12 +89,17 @@ def check_success(self, cmd, result) -> bool:
         last_line = ansi_escape.sub("", last_line)
         return got_root(self.conn.hostname, last_line)
 
+    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()
-        result, task_successful = self.run_command(cmd, message_id)
+
+        cmds = self.postprocess_commands(cmd)
+        for cmd in cmds:
+            result, task_successful = self.run_command(cmd, message_id)
 
         # maybe move the 'got root' detection here?
         # TODO: also can I use llm-as-judge for that? or do I have to do this

src/hackingBuddyGPT/usecases/__init__.py

@@ -1,8 +1,7 @@
-from .examples import *
-from .privesc import *
 from .web import *
 from .web_api_testing import *
 from .viewer import *
 from .rag import *
 from .minimal_linux_privesc import *
+from .call_usecase_from_usecase import *
 from .linux_privesc import *

…hackingBuddyGPT/usecases/examples/lse.py …PT/usecases/call_usecase_from_usecase.pysrc/hackingBuddyGPT/usecases/examples/lse.py renamed to src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py src/hackingBuddyGPT/usecases/examples/lse.py renamed to src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py

@@ -1,16 +1,22 @@
-import pathlib
-
 from mako.template import Template
 
 from hackingBuddyGPT.capabilities import SSHRunCommand
 from hackingBuddyGPT.usecases.base import UseCase, use_case
-from hackingBuddyGPT.usecases.privesc.linux import LinuxPrivesc, LinuxPrivescUseCase
 from hackingBuddyGPT.utils import SSHConnection
 from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
 
-template_dir = pathlib.Path(__file__).parent
-template_lse = Template(filename=str(template_dir / "get_hint_from_lse.txt"))
+from .linux_privesc import PrivEscLinux
+
+template_lse = Template("""
+Create a list of up to ${number} attack classes that you would try on a linux system
+(to achieve root level privileges) given the following output:
 
+~~~ bash
+${lse_output}
+~~~
+
+only output the list of attack classes, for each attack class only output a single
+short sentence.""")
 
 @use_case("Linux Privilege Escalation using lse.sh for initial guidance")
 class ExPrivEscLinuxLSEUseCase(UseCase):
@@ -23,9 +29,6 @@ class ExPrivEscLinuxLSEUseCase(UseCase):
 
     _got_root: bool = False
 
-    # use either an use-case or an agent to perform the privesc
-    use_use_case: bool = False
-
     # simple helper that uses lse.sh to get hints from the system
     def call_lse_against_host(self):
         self.log.console.print("[green]performing initial enumeration with lse.sh")
@@ -43,65 +46,32 @@ def call_lse_against_host(self):
     def get_name(self) -> str:
         return self.__class__.__name__
 
-    def run(self):
+    def run(self, configuration={}):
         # get the hints through running LSE on the target system
         hints = self.call_lse_against_host()
         turns_per_hint = int(self.max_turns / len(hints))
 
         # now try to escalate privileges using the hints
         for hint in hints:
-            if self.use_use_case:
-                self.log.console.print("[yellow]Calling a use-case to perform the privilege escalation")
-                result = self.run_using_usecases(hint, turns_per_hint)
-            else:
-                self.log.console.print("[yellow]Calling an agent to perform the privilege escalation")
-                result = self.run_using_agent(hint, turns_per_hint)
+            self.log.console.print("[yellow]Calling a use-case to perform the privilege escalation")
+            result = self.run_using_usecases(hint, turns_per_hint)
 
             if result is True:
                 self.log.console.print("[green]Got root!")
                 return True
 
     def run_using_usecases(self, hint, turns_per_hint):
         # TODO: init usecase
-        linux_privesc = LinuxPrivescUseCase(
-            agent=LinuxPrivesc(
-                conn=self.conn,
-                enable_explanation=self.enable_explanation,
-                enable_update_state=self.enable_update_state,
-                disable_history=self.disable_history,
-                llm=self.llm,
-                hint=hint,
-            ),
-            max_turns=turns_per_hint,
-            log=self.log,
-        )
-        linux_privesc.init(self.configuration)
-        return linux_privesc.run()
-
-    def run_using_agent(self, hint, turns_per_hint):
-        # init agent
-        agent = LinuxPrivesc(
+        linux_privesc = PrivEscLinux(
             conn=self.conn,
-            llm=self.llm,
-            hint=hint,
             enable_explanation=self.enable_explanation,
             enable_update_state=self.enable_update_state,
             disable_history=self.disable_history,
+            llm=self.llm,
+            hints=f"hint:{hint}",
+            max_turns=turns_per_hint,
+            log=self.log,
         )
-        agent.log = self.log
-        agent.init()
-
-        # perform the privilege escalation
-        agent.before_run()
-        turn = 1
-        got_root = False
-        while turn <= turns_per_hint and not got_root:
-            self.log.console.log(f"[yellow]Starting turn {turn} of {turns_per_hint}")
-
-            if agent.perform_round(turn) is True:
-                got_root = True
-            turn += 1
-
-        # cleanup and finish
-        agent.after_run()
-        return got_root
+
+        linux_privesc.init()
+        return linux_privesc.run({})