introduct before_run/after_run hooks and use them within the hintfile use-case

Author: andreashappe

src/hackingBuddyGPT/usecases/agents.py

@@ -21,6 +21,12 @@ class Agent(ABC):
     def init(self):
         pass
 
+    def before_run(self):
+        pass
+
+    def after_run(self):
+        pass
+
     # callback
     @abstractmethod
     def perform_round(self, turn: int) -> bool:

src/hackingBuddyGPT/usecases/base.py

@@ -74,7 +74,16 @@ class AutonomousUseCase(UseCase, abc.ABC):
     def perform_round(self, turn: int):
         pass
 
+    def before_run(self):
+        pass
+
+    def after_run(self):
+        pass
+
     def run(self):
+
+        self.before_run()
+
         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}")
@@ -85,6 +94,8 @@ def run(self):
             self._log.log_db.commit()
             turn += 1
 
+        self.after_run()
+
         # write the final result to the database and console
         if self._got_root:
             self._log.log_db.run_was_success(self._run_id, turn)
@@ -145,6 +156,12 @@ def init(self):
 
             def get_name(self) -> str:
                 return self.__class__.__name__
+            
+            def before_run(self):
+                return self.agent.before_run()
+            
+            def after_run(self):
+                return self.agent.after_run()
 
             def perform_round(self, turn: int):
                 return self.agent.perform_round(turn)

src/hackingBuddyGPT/usecases/privesc/common.py

@@ -35,6 +35,7 @@ class Privesc(Agent):
     def init(self):
         super().init()
 
+    def before_run(self):
         if self.hint != "":
             self._log.console.print(f"[bold green]Using the following hint: '{self.hint}'")
 

src/hackingBuddyGPT/usecases/privesc/linux.py

@@ -32,22 +32,20 @@ class LinuxPrivescWithHintFileUseCase(AutonomousAgentUseCase[LinuxPrivesc]):
 
     def init(self):
         super().init()
-
         self.agent.hint = self.read_hint()
 
     # simple helper that reads the hints file and returns the hint
     # for the current machine (test-case)
     def read_hint(self):
-        if self.hints != "":
-            try:
-                with open(self.hints, "r") as hint_file:
-                    hints = json.load(hint_file)
-                    if self.agent.conn.hostname in hints:
-                        return hints[self.agent.conn.hostname]
-            except:
-                self._log.console.print("[yellow]Was not able to load hint file")
-        else:
-            self._log.console.print("[yellow]calling the hintfile use-case without a hint file?")
+        try:
+            with open(self.hints, "r") as hint_file:
+                hints = json.load(hint_file)
+                if self.agent.conn.hostname in hints:
+                    return hints[self.agent.conn.hostname]
+        except FileNotFoundError:
+            self._log.console.print("[yellow]Hint file not found")
+        except Exception as e:
+            self._log.console.print("[yellow]Hint file could not loaded:", str(e))
         return ""
 
 
@@ -58,7 +56,6 @@ class LinuxPrivescWithLSEUseCase(AutonomousAgentUseCase[LinuxPrivesc]):
 
     def init(self):
         super().init()
-
         self._hints = self.read_hint().splitlines()
         self._turns_per_hint = int(self.max_turns / len(self._hints))