Merge pull request #73 from ipa-lab/explorative_refactoring

Explorative refactoring

Author: andreashappe

src/hackingBuddyGPT/usecases/agents.py

@@ -4,18 +4,27 @@
 from rich.panel import Panel
 from typing import Dict
 
+from hackingBuddyGPT.usecases.base import Logger
 from hackingBuddyGPT.utils import llm_util
-
 from hackingBuddyGPT.capabilities.capability import Capability, capabilities_to_simple_text_handler
-from .common_patterns import RoundBasedUseCase
+from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
+
 
 @dataclass
-class Agent(RoundBasedUseCase, ABC):
+class Agent(ABC):
     _capabilities: Dict[str, Capability] = field(default_factory=dict)
     _default_capability: Capability = None
+    _log: Logger = None
+
+    llm: OpenAIConnection = None
 
     def init(self):
-        super().init()
+        pass
+
+    # callback
+    @abstractmethod
+    def perform_round(self, turn: int) -> bool:
+        pass
 
     def add_capability(self, cap: Capability, default: bool = False):
         self._capabilities[cap.get_name()] = cap
@@ -29,6 +38,7 @@ def get_capability_block(self) -> str:
         capability_descriptions, _parser = capabilities_to_simple_text_handler(self._capabilities)
         return "You can either\n\n" + "\n".join(f"- {description}" for description in capability_descriptions.values())
 
+
 @dataclass
 class AgentWorldview(ABC):
 
@@ -40,6 +50,7 @@ def to_template(self):
     def update(self, capability, cmd, result):
         pass
 
+
 class TemplatedAgent(Agent):
 
     _state: AgentWorldview = None
@@ -59,7 +70,7 @@ def set_template(self, template:str):
     def perform_round(self, turn:int) -> bool:
         got_root : bool = False
 
-        with self.console.status("[bold green]Asking LLM for a new command..."):
+        with self._log.console.status("[bold green]Asking LLM for a new command..."):
             # TODO output/log state
             options = self._state.to_template()
             options.update({
@@ -70,16 +81,16 @@ def perform_round(self, turn:int) -> bool:
             answer = self.llm.get_response(self._template, **options)
             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:"))
+        with self._log.console.status("[bold green]Executing that command..."):
+                self._log.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._log.log_db.add_log_query(self._log.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}"))
+        self._log.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
 
         # if we got root, we can stop the loop
         return got_root

src/hackingBuddyGPT/usecases/base.py

@@ -1,10 +1,24 @@
 import abc
 import argparse
-from dataclasses import dataclass, field
+import typing
+from dataclasses import dataclass
+from rich.panel import Panel
 from typing import Dict, Type
 
-from hackingBuddyGPT.utils.configurable import ParameterDefinitions, build_parser, get_arguments, get_class_parameters
+from hackingBuddyGPT.utils.configurable import ParameterDefinitions, build_parser, get_arguments, get_class_parameters, transparent
+from hackingBuddyGPT.utils.console.console import Console
+from hackingBuddyGPT.utils.db_storage.db_storage import DbStorage
 
+
+@dataclass
+class Logger:
+    log_db: DbStorage
+    console: Console
+    tag: str = ""
+    run_id: int = 0
+
+
+@dataclass
 class UseCase(abc.ABC):
     """
     A UseCase is the combination of tools and capabilities to solve a specific problem.
@@ -16,13 +30,21 @@ class UseCase(abc.ABC):
     so that they can be automatically discovered and run from the command line.
     """
 
+    log_db: DbStorage
+    console: Console
+    tag: str = ""
+
+    _run_id: int = 0
+    _log: Logger = None
+
     def init(self):
         """
         The init method is called before the run method. It is used to initialize the UseCase, and can be used to
         perform any dynamic setup that is needed before the run method is called. One of the most common use cases is
         setting up the llm capabilities from the tools that were injected.
         """
-        pass
+        self._run_id = self.log_db.create_new_run(self.get_name(), self.tag)
+        self._log = Logger(self.log_db, self.console, self.tag, self._run_id)
 
     @abc.abstractmethod
     def run(self):
@@ -33,6 +55,46 @@ def run(self):
         """
         pass
 
+    @abc.abstractmethod
+    def get_name(self) -> str:
+        """
+        This method should return the name of the use case. It is used for logging and debugging purposes.
+        """
+        pass
+
+
+# this runs the main loop for a bounded amount of turns or until root was achieved
+@dataclass
+class AutonomousUseCase(UseCase, abc.ABC):
+    max_turns: int = 10
+
+    _got_root: bool = False
+
+    @abc.abstractmethod
+    def perform_round(self, turn: int):
+        pass
+
+    def run(self):
+        turn = 1
+        while turn <= self.max_turns and not self._got_root:
+            self._log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}")
+
+            self._got_root = self.perform_round(turn)
+
+            # finish turn and commit logs to storage
+            self._log.log_db.commit()
+            turn += 1
+
+        # write the final result to the database and console
+        if self._got_root:
+            self._log.log_db.run_was_success(self._run_id, turn)
+            self._log.console.print(Panel("[bold green]Got Root!", title="Run finished"))
+        else:
+            self._log.log_db.run_was_failure(self._run_id, turn)
+            self._log.console.print(Panel("[green]maximum turn number reached", title="Run finished"))
+
+        return self._got_root
+
 
 @dataclass
 class _WrappedUseCase:
@@ -56,17 +118,55 @@ def __call__(self, args: argparse.Namespace):
 use_cases: Dict[str, _WrappedUseCase] = dict()
 
 
-def use_case(name: str, desc: str):
-    """
-    By wrapping a UseCase with this decorator, it will be automatically discoverable and can be run from the command
-    line.
-    """
+T = typing.TypeVar("T")
+
+
+class AutonomousAgentUseCase(AutonomousUseCase, typing.Generic[T]):
+    agent: T = None
+
+    def perform_round(self, turn: int):
+        raise ValueError("Do not use AutonomousAgentUseCase without supplying an agent type as generic")
+
+    def get_name(self) -> str:
+        raise ValueError("Do not use AutonomousAgentUseCase without supplying an agent type as generic")
+
+    @classmethod
+    def __class_getitem__(cls, item):
+        item = dataclass(item)
+        item.__parameters__ = get_class_parameters(item)
 
-    def inner(cls: Type[UseCase]):
+        class AutonomousAgentUseCase(AutonomousUseCase):
+            agent: transparent(item) = None
+
+            def init(self):
+                super().init()
+                self.agent._log = self._log
+                self.agent.init()
+
+            def get_name(self) -> str:
+                return self.__class__.__name__
+
+            def perform_round(self, turn: int):
+                return self.agent.perform_round(turn)
+
+        constructed_class = dataclass(AutonomousAgentUseCase)
+
+        return constructed_class
+
+
+def use_case(description):
+    def inner(cls):
+        name = cls.__name__.removesuffix("UseCase")
         if name in use_cases:
             raise IndexError(f"Use case with name {name} already exists")
-        use_cases[name] = _WrappedUseCase(name, desc, cls, get_class_parameters(cls, name))
+        use_cases[name] = _WrappedUseCase(name, description, cls, get_class_parameters(cls))
+    return inner
 
-        return cls
 
-    return inner
+def register_use_case(name: str, description: str, use_case: Type[UseCase]):
+    """
+    This function is used to register a UseCase that was created manually, and not through the use_case decorator.
+    """
+    if name in use_cases:
+        raise IndexError(f"Use case with name {name} already exists")
+    use_cases[name] = _WrappedUseCase(name, description, use_case, get_class_parameters(use_case))

src/hackingBuddyGPT/usecases/common_patterns.py

@@ -1,23 +1,20 @@
 import pathlib
-from dataclasses import dataclass, field
 from mako.template import Template
 from rich.panel import Panel
 
 from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential
 from hackingBuddyGPT.utils import SSHConnection, llm_util
-from hackingBuddyGPT.usecases.base import use_case
+from hackingBuddyGPT.usecases.base import use_case, AutonomousAgentUseCase
 from hackingBuddyGPT.usecases.agents import Agent
 from hackingBuddyGPT.utils.cli_history import SlidingCliHistory
 
 template_dir = pathlib.Path(__file__).parent
 template_next_cmd = Template(filename=str(template_dir / "next_cmd.txt"))
 
-@use_case("minimal_linux_privesc", "Showcase Minimal Linux Priv-Escalation")
-@dataclass
+
 class MinimalLinuxPrivesc(Agent):
 
     conn: SSHConnection = None
-    
     _sliding_history: SlidingCliHistory = None
 
     def init(self):
@@ -27,25 +24,30 @@ def init(self):
         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
+    def perform_round(self, turn: int) -> bool:
+        got_root: bool = False
 
-        with self.console.status("[bold green]Asking LLM for a new command..."):
+        with self._log.console.status("[bold green]Asking LLM for a new command..."):
             # get as much history as fits into the target context size
             history = self._sliding_history.get_history(self.llm.context_size - llm_util.SAFETY_MARGIN - self._template_size)
 
             # get the next command from the LLM
             answer = self.llm.get_response(template_next_cmd, capabilities=self.get_capability_block(), history=history, conn=self.conn)
             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)
+        with self._log.console.status("[bold green]Executing that command..."):
+            self._log.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:"))
+            result, got_root = self.get_capability(cmd.split(" ", 1)[0])(cmd)
 
         # log and output the command and its result
-        self.log_db.add_log_query(self._run_id, turn, cmd, result, answer)
+        self._log.log_db.add_log_query(self._log.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.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
 
         # if we got root, we can stop the loop
         return got_root
+
+
+@use_case("Showcase Minimal Linux Priv-Escalation")
+class MinimalLinuxPrivescUseCase(AutonomousAgentUseCase[MinimalLinuxPrivesc]):
+    pass