add more test-cases (also test github integration)

.gitignore

@@ -10,3 +10,4 @@ __pycache__/
 src/hackingBuddyGPT.egg-info/
 build/
 dist/
+.coverage

pyproject.toml

@@ -14,7 +14,7 @@ maintainers = [
 description = "Helping Ethical Hackers use LLMs in 50 lines of code"
 readme = "README.md"
 keywords = ["hacking", "pen-testing", "LLM", "AI", "agent"]
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 version = "0.3.0"
 license = { file = "LICENSE" }
 classifiers = [

src/hackingBuddyGPT/usecases/privesc/common.py

@@ -16,6 +16,7 @@
 template_state = Template(filename=str(template_dir / "update_state.txt"))
 template_lse = Template(filename=str(template_dir / "get_hint_from_lse.txt"))
 
+
 @dataclass
 class Privesc(Agent):
 
@@ -24,20 +25,20 @@ class Privesc(Agent):
     enable_update_state: bool = False
     disable_history: bool = False
     hint: str = ""
-    
+
     _sliding_history: SlidingCliHistory = None
     _state: str = ""
     _capabilities: Dict[str, Capability] = field(default_factory=dict)
-    _template_params : Dict[str, Any] = field(default_factory=dict)
-    _max_history_size : int = 0
+    _template_params: Dict[str, Any] = field(default_factory=dict)
+    _max_history_size: int = 0
 
     def init(self):
         super().init()
 
     def setup(self):
         if self.hint != "":
             self.console.print(f"[bold green]Using the following hint: '{self.hint}'")
-        
+
         if self.disable_history is False:
             self._sliding_history = SlidingCliHistory(self.llm)
 
@@ -53,8 +54,8 @@ def setup(self):
         template_size = self.llm.count_tokens(template_next_cmd.source)
         self._max_history_size = self.llm.context_size - llm_util.SAFETY_MARGIN - template_size
 
-    def perform_round(self, turn:int) -> bool:
-        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..."):
             answer = self.get_next_command()
@@ -65,7 +66,7 @@ def perform_round(self, turn:int) -> bool:
             _capability_descriptions, parser = capabilities_to_simple_text_handler(self._capabilities, default_capability=self._default_capability)
             success, *output = parser(cmd)
             if not success:
-                self.console.print(Panel(output[0], title=f"[bold red]Error parsing command:"))
+                self.console.print(Panel(output[0], title="[bold red]Error parsing command:"))
                 return False
 
             assert(len(output) == 1)
@@ -140,4 +141,4 @@ def update_state(self, cmd, result):
 
         result = self.llm.get_response(template_state, cmd=cmd, resp=result, facts=self._state)
         self._state = result.result
-        return result
+        return result

src/hackingBuddyGPT/utils/cli_history.py

@@ -14,5 +14,5 @@ def add_command(self, cmd: str, output: str):
         self.sliding_history += f"$ {cmd}\n{output}"
         self.sliding_history = trim_result_front(self.model, self.maximum_target_size, self.sliding_history)
 
-    def get_history(self, target_size:int) -> str:
+    def get_history(self, target_size: int) -> str:
         return trim_result_front(self.model, min(self.maximum_target_size, target_size), self.sliding_history)

src/hackingBuddyGPT/utils/shell_root_detection.py

@@ -5,10 +5,11 @@
     re.compile("^bash-[0-9]+.[0-9]# $")
 ]
 
-def got_root(hostname:str, output:str) -> bool:
+
+def got_root(hostname: str, output: str) -> bool:
     for i in GOT_ROOT_REGEXPs:
         if i.fullmatch(output):
             return True
     if output.startswith(f'root@{hostname}:'):
         return True
-    return False
+    return False

tests/integration_minimal_test.py

@@ -1,5 +1,7 @@
 
 from typing import Tuple
+from hackingBuddyGPT.usecases.minimal.agent import MinimalLinuxPrivesc
+from hackingBuddyGPT.usecases.minimal.agent_with_state import MinimalLinuxTemplatedPrivesc
 from hackingBuddyGPT.usecases.privesc.linux import LinuxPrivesc
 from hackingBuddyGPT.utils.console.console import Console
 from hackingBuddyGPT.utils.db_storage.db_storage import DbStorage
@@ -63,7 +65,7 @@ def get_response(self, prompt, *, capabilities=None, **kwargs) -> LLMResult:
     def encode(self, query) -> list[int]:
         return [0]
 
-def test_minimal_linuxprives():
+def test_linuxprivesc():
 
     conn = FakeSSHConnection()
     llm = FakeLLM()
@@ -85,7 +87,49 @@ def test_minimal_linuxprives():
     )
 
     priv_esc.init()
+    result = priv_esc.run()
+    assert result is True
+
+def test_minimal_agent():
+
+    conn = FakeSSHConnection()
+    llm = FakeLLM()
+    log_db = DbStorage(':memory:')
+    console = Console()
+
+    log_db.init()
 
+    priv_esc = MinimalLinuxPrivesc(
+        conn=conn,
+        log_db = log_db,
+        console = console,
+        llm = llm,
+        tag = 'integration_test_linuxprivesc',
+        max_turns = len(llm.responses)
+    )
+
+    priv_esc.init()
     result = priv_esc.run()
+    assert result is True
 
+def test_minimal_agent_state():
+
+    conn = FakeSSHConnection()
+    llm = FakeLLM()
+    log_db = DbStorage(':memory:')
+    console = Console()
+
+    log_db.init()
+
+    priv_esc = MinimalLinuxTemplatedPrivesc(
+        conn=conn,
+        log_db = log_db,
+        console = console,
+        llm = llm,
+        tag = 'integration_test_linuxprivesc',
+        max_turns = len(llm.responses)
+    )
+
+    priv_esc.init()
+    result = priv_esc.run()
     assert result is True