Merge pull request #79 from ipa-lab/log_infrastructure

Semantic logging #75

Author: andreashappe

.gitignore

@@ -1,5 +1,6 @@
 .env
 venv/
+.venv/
 __pycache__/
 *.swp
 *.log

README.md

@@ -12,7 +12,7 @@ If you want to use hackingBuddyGPT and need help selecting the best LLM for your
 
 ## hackingBuddyGPT in the News
 
-- **upcoming** 2024-11-20: [Manuel Reinsperger](https://www.github.com/neverbolt) will present hackingBuddyGPT at the [European Symposium on Security and Artificial Intelligence (ESSAI)](https://essai-conference.eu/) 
+- 2024-11-20: [Manuel Reinsperger](https://www.github.com/neverbolt) presented hackingBuddyGPT at the [European Symposium on Security and Artificial Intelligence (ESSAI)](https://essai-conference.eu/) 
 - 2024-07-26: The [GitHub Accelerator Showcase](https://github.blog/open-source/maintainers/github-accelerator-showcase-celebrating-our-second-cohort-and-whats-next/) features hackingBuddyGPT
 - 2024-07-24: [Juergen](https://github.com/citostyle) speaks at [Open Source + mezcal night @ GitHub HQ](https://lu.ma/bx120myg)
 - 2024-05-23: hackingBuddyGPT is part of [GitHub Accelerator 2024](https://github.blog/news-insights/company-news/2024-github-accelerator-meet-the-11-projects-shaping-open-source-ai/)
@@ -82,38 +82,38 @@ template_next_cmd = Template(filename=str(template_dir / "next_cmd.txt"))
 
 
 class MinimalLinuxPrivesc(Agent):
-
     conn: SSHConnection = None
+
     _sliding_history: SlidingCliHistory = None
+    _max_history_size: int = 0
 
     def init(self):
         super().init()
+
         self._sliding_history = SlidingCliHistory(self.llm)
+        self._max_history_size = self.llm.context_size - llm_util.SAFETY_MARGIN - self.llm.count_tokens(template_next_cmd.source)
+
         self.add_capability(SSHRunCommand(conn=self.conn), default=True)
         self.add_capability(SSHTestCredential(conn=self.conn))
-        self._template_size = self.llm.count_tokens(template_next_cmd.source)
 
-    def perform_round(self, turn: int) -> bool:
-        got_root: bool = False
+    @log_conversation("Asking LLM for a new command...")
+    def perform_round(self, turn: int, log: Logger) -> bool:
+        # get as much history as fits into the target context size
+        history = self._sliding_history.get_history(self._max_history_size)
 
-        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)
+        message_id = log.call_response(answer)
 
-            # 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)
+        # clean the command, load and execute it
+        cmd = llm_util.cmd_output_fixer(answer.result)
+        capability, arguments = cmd.split(" ", 1)
+        result, got_root = self.run_capability(message_id, "0", capability, arguments, calling_mode=CapabilityCallingMode.Direct, log=log)
 
-        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.log_db.add_log_query(self._log.run_id, turn, cmd, result, answer)
+        # store the results in our local history
         self._sliding_history.add_command(cmd, result)
-        self._log.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
 
-        # if we got root, we can stop the loop
+        # signal if we were successful in our task
         return got_root
 
 
@@ -306,6 +306,22 @@ Mac, Docker Desktop and Gemini-OpenAI-Proxy:
 
 * See https://github.com/ipa-lab/hackingBuddyGPT/blob/main/MAC.md
 
+## Beta Features
+
+### Viewer
+
+The viewer is a simple web-based tool to view the results of hackingBuddyGPT runs. It is currently in beta and can be started with:
+
+```bash
+$ hackingBuddyGPT Viewer
+```
+
+This will start a webserver on `http://localhost:4444` that can be accessed with a web browser.
+
+To log to this central viewer, you currently need to change the `GlobalLogger` definition in [./src/hackingBuddyGPT/utils/logging.py](src/hackingBuddyGPT/utils/logging.py) to `GlobalRemoteLogger`.
+
+This feature is not fully tested yet and therefore is not recommended to be exposed to the internet!
+
 ## Publications about hackingBuddyGPT
 
 Given our background in academia, we have authored papers that lay the groundwork and report on our efforts:

pyproject.toml

@@ -26,19 +26,25 @@ classifiers = [
 	"Development Status :: 4 - Beta",
 ]
 dependencies = [
-	'fabric == 3.2.2',
-	'Mako == 1.3.2',
-	'requests == 2.32.0',
-	'rich == 13.7.1',
-	'tiktoken == 0.8.0',
-	'instructor == 1.3.5',
-	'PyYAML == 6.0.1',
-	'python-dotenv == 1.0.1',
-	'pypsexec == 0.3.0',
-	'pydantic == 2.8.2',
-	'openai == 1.28.0',
-	'BeautifulSoup4',
-	'nltk'
+    'fabric == 3.2.2',
+    'Mako == 1.3.2',
+    'requests == 2.32.0',
+    'rich == 13.7.1',
+    'tiktoken == 0.8.0',
+    'instructor == 1.3.5',
+    'PyYAML == 6.0.1',
+    'python-dotenv == 1.0.1',
+    'pypsexec == 0.3.0',
+    'pydantic == 2.8.2',
+    'openai == 1.28.0',
+    'BeautifulSoup4',
+    'nltk',
+    'fastapi == 0.114.0',
+    'fastapi-utils == 0.7.0',
+    'jinja2 == 3.1.4',
+    'uvicorn[standard] == 0.30.6',
+    'dataclasses_json == 0.6.7',
+    'websockets == 13.1',
 ]
 
 [project.urls]
@@ -56,14 +62,9 @@ where = ["src"]
 
 [tool.pytest.ini_options]
 pythonpath = "src"
-addopts = [
-	"--import-mode=importlib",
-]
+addopts = ["--import-mode=importlib"]
 [project.optional-dependencies]
-testing = [
-    'pytest',
-    'pytest-mock'
-]
+testing = ['pytest', 'pytest-mock']
 dev = [
 	'ruff',
 ]

src/hackingBuddyGPT/capabilities/capability.py

@@ -38,14 +38,14 @@ def get_name(self) -> str:
     def __call__(self, *args, **kwargs):
         """
         The actual execution of a capability, please make sure, that the parameters and return type of your
-        implementation are well typed, as this will make it easier to support full function calling soon.
+        implementation are well typed, as this is used to properly support function calling.
         """
         pass
 
     def to_model(self) -> BaseModel:
         """
         Converts the parameters of the `__call__` function of the capability to a pydantic model, that can be used to
-        interface with an LLM using eg instructor or the openAI function calling API.
+        interface with an LLM using eg the openAI function calling API.
         The model will have the same name as the capability class and will have the same fields as the `__call__`,
         the `__call__` method can then be accessed by calling the `execute` method of the model.
         """

src/hackingBuddyGPT/capabilities/http_request.py

@@ -70,7 +70,7 @@ def __call__(
             url = self.host + ("" if path.startswith("/") else "/") + path + ("?{query}" if query else "")
             return f"Could not request '{url}': {e}"
 
-        headers = "\r\n".join(f"{k}: {v}" for k, v in resp.headers.items())
+        response_headers = "\r\n".join(f"{k}: {v}" for k, v in resp.headers.items())
 
         # turn the response into "plain text format" for responding to the prompt
-        return f"HTTP/1.1 {resp.status_code} {resp.reason}\r\n{headers}\r\n\r\n{resp.text}"
+        return f"HTTP/1.1 {resp.status_code} {resp.reason}\r\n{response_headers}\r\n\r\n{resp.text}"

src/hackingBuddyGPT/capabilities/submit_http_method.py

@@ -1,10 +1,8 @@
 import base64
-import inspect
 from dataclasses import dataclass, field
 from typing import Callable, Dict, Literal, Optional, Set
 
 import requests
-from pydantic import BaseModel, create_model
 
 from . import Capability
 
@@ -23,25 +21,6 @@ class SubmitHTTPMethod(Capability):
     def describe(self) -> str:
         return f"Submits a HTTP methods, which is the goal of all activities. The HTTP methods format is {self.http_methods_format}. If you find a flag following the that format, that you have not yet submitted, then the most important first step is to submit it, ignoring all other possibilities of further action"
 
-    def to_model(self) -> BaseModel:
-        """
-        Converts the parameters of the `__call__` function of the capability to a pydantic model, that can be used to
-        interface with an LLM using eg instructor or the openAI function calling API.
-        The model will have the same name as the capability class and will have the same fields as the `__call__`,
-        the `__call__` method can then be accessed by calling the `execute` method of the model.
-        """
-        sig = inspect.signature(self.__call__)
-        fields = {param: (param_info.annotation, ...) for param, param_info in sig.parameters.items()}
-        model_type = create_model(self.__class__.__name__, __doc__=self.describe(), **fields)
-
-        def execute(model):
-            m = model.dict()
-            return self(**m)
-
-        model_type.execute = execute
-
-        return model_type
-
     def __call__(
         self,
         method: Literal["GET", "HEAD", "POST", "PUT", "DELETE", "OPTION", "PATCH"],

src/hackingBuddyGPT/cli/stats.py

@@ -11,8 +11,9 @@ def main():
         use_case.build_parser(subparser.add_parser(name=name, help=use_case.description))
 
     parsed = parser.parse_args(sys.argv[1:])
+    configuration = {k: v for k, v in vars(parsed).items() if k not in ("use_case", "parser_state")}
     instance = parsed.use_case(parsed)
-    instance.init()
+    instance.init(configuration=configuration)
     instance.run()