Merge branch 'development' into new-branch-1

.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
 
 
@@ -219,6 +219,22 @@ $ python src/hackingBuddyGPT/cli/wintermute.py LinuxPrivesc --llm.api_key=sk...C
 $ pip install '.[testing]'
 ```
 
+## 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:

publish_notes.md

@@ -0,0 +1,34 @@
+# how to publish to pypi
+
+## start with testing if the project builds and tag the version
+
+```bash
+python -m venv venv
+source venv/bin/activate
+pip install -e .
+pytest
+git tag v0.3.0
+git push origin v0.3.0
+```
+
+## build and new package
+
+(according to https://packaging.python.org/en/latest/tutorials/packaging-projects/)
+
+```bash
+pip install build twine
+python3 -m build
+vi ~/.pypirc
+twine check dist/*
+```
+
+Now, for next time.. test install the package in a new vanilla environment, then..
+
+```bash
+twine upload dist/*
+```
+
+## repo todos
+
+- rebase development upon main
+- bump the pyproject version number to a new `-dev` 

pyproject.toml

@@ -4,18 +4,20 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hackingBuddyGPT"
+# original author was Andreas Happe, for an up-to-date list see
+# https://github.com/ipa-lab/hackingBuddyGPT/graphs/contributors
 authors = [
-	{ name = "Andreas Happe", email = "andreas@offensive.one" }
+	{ name = "HackingBuddyGPT maintainers", email = "maintainers@hackingbuddy.ai" }
 ]
 maintainers = [
 	{ name = "Andreas Happe", email = "andreas@offensive.one" },
-	{ name = "Juergen Cito", email = "juergen.cito@tuwiena.c.at" }
+	{ name = "Juergen Cito", email = "juergen.cito@tuwien.ac.at" }
 ]
 description = "Helping Ethical Hackers use LLMs in 50 lines of code"
 readme = "README.md"
 keywords = ["hacking", "pen-testing", "LLM", "AI", "agent"]
 requires-python = ">=3.10"
-version = "0.3.1"
+version = "0.4.0-dev"
 license = { file = "LICENSE" }
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -24,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]
@@ -54,15 +62,20 @@ 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',
 ]
 
 [project.scripts]
 wintermute = "hackingBuddyGPT.cli.wintermute:main"
 hackingBuddyGPT = "hackingBuddyGPT.cli.wintermute:main"
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+select = ["E", "F", "B", "I"]
+ignore = ["E501", "F401", "F403"]

src/hackingBuddyGPT/capabilities/__init__.py

@@ -1,5 +1,13 @@
 from .capability import Capability
-from .psexec_test_credential import PSExecTestCredential
 from .psexec_run_command import PSExecRunCommand
+from .psexec_test_credential import PSExecTestCredential
 from .ssh_run_command import SSHRunCommand
-from .ssh_test_credential import SSHTestCredential
+from .ssh_test_credential import SSHTestCredential
+
+__all__ = [
+    "Capability",
+    "PSExecRunCommand",
+    "PSExecTestCredential",
+    "SSHRunCommand",
+    "SSHTestCredential",
+]

src/hackingBuddyGPT/capabilities/capability.py

@@ -1,11 +1,11 @@
 import abc
 import inspect
-from typing import Union, Type, Dict, Callable, Any, Iterable
+from typing import Any, Callable, Dict, Iterable, Type, Union
 
 import openai
 from openai.types.chat import ChatCompletionToolParam
 from openai.types.chat.completion_create_params import Function
-from pydantic import create_model, BaseModel
+from pydantic import BaseModel, create_model
 
 
 class Capability(abc.ABC):
@@ -18,12 +18,13 @@ class Capability(abc.ABC):
     At the moment, this is not yet a very powerful class, but in the near-term future, this will provide an automated
     way of providing a json schema for the capabilities, which can then be used for function-calling LLMs.
     """
+
     @abc.abstractmethod
     def describe(self) -> str:
         """
         describe should return a string that describes the capability. This is used to generate the help text for the
         LLM.
- 
+
         This is a method and not just a simple property on purpose (though it could become a @property in the future, if
         we don't need the name parameter anymore), so that it can template in some of the capabilities parameters into
         the description.
@@ -37,23 +38,30 @@ 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.
         """
         sig = inspect.signature(self.__call__)
-        fields = {param: (param_info.annotation, param_info.default if param_info.default is not inspect._empty else ...) for param, param_info in sig.parameters.items()}
+        fields = {
+            param: (
+                param_info.annotation,
+                param_info.default if param_info.default is not inspect._empty else ...,
+            )
+            for param, param_info in sig.parameters.items()
+        }
         model_type = create_model(self.__class__.__name__, __doc__=self.describe(), **fields)
 
         def execute(model):
             return self(**model.dict())
+
         model_type.execute = execute
 
         return model_type
@@ -76,6 +84,7 @@ def capabilities_to_action_model(capabilities: Dict[str, Capability]) -> Type[Ac
     This allows the LLM to define an action to be used, which can then simply be called using the `execute` function on
     the model returned from here.
     """
+
     class Model(Action):
         action: Union[tuple([capability.to_model() for capability in capabilities.values()])]
 
@@ -86,7 +95,11 @@ class Model(Action):
 SimpleTextHandler = Callable[[str], SimpleTextHandlerResult]
 
 
-def capabilities_to_simple_text_handler(capabilities: Dict[str, Capability], default_capability: Capability = None, include_description: bool = True) -> tuple[Dict[str, str], SimpleTextHandler]:
+def capabilities_to_simple_text_handler(
+    capabilities: Dict[str, Capability],
+    default_capability: Capability = None,
+    include_description: bool = True,
+) -> tuple[Dict[str, str], SimpleTextHandler]:
     """
     This function generates a simple text handler from a set of capabilities.
     It is to be used when no function calling is available, and structured output is not to be trusted, which is why it
@@ -97,12 +110,16 @@ def capabilities_to_simple_text_handler(capabilities: Dict[str, Capability], def
     whether the parsing was successful, the second return value is a tuple containing the capability name, the parameters
     as a string and the result of the capability execution.
     """
+
     def get_simple_fields(func, name) -> Dict[str, Type]:
         sig = inspect.signature(func)
         fields = {param: param_info.annotation for param, param_info in sig.parameters.items()}
         for param, param_type in fields.items():
             if param_type not in (str, int, float, bool):
-                raise ValueError(f"The command {name} is not compatible with this calling convention (this is not a LLM error, but rather a problem with the capability itself, the parameter {param} is {param_type} and not a simple type (str, int, float, bool))")
+                raise ValueError(
+                    f"The command {name} is not compatible with this calling convention (this is not a LLM error,"
+                    f"but rather a problem with the capability itself, the parameter {param} is {param_type} and not a simple type (str, int, float, bool))"
+                )
         return fields
 
     def parse_params(fields, params) -> tuple[bool, Union[str, Dict[str, Any]]]:
@@ -169,13 +186,14 @@ def default_capability_parser(text: str) -> SimpleTextHandlerResult:
 
             return True, (capability_name, params, default_capability(**parsing_result))
 
-
         resolved_parser = default_capability_parser
 
     return capability_descriptions, resolved_parser
 
 
-def capabilities_to_functions(capabilities: Dict[str, Capability]) -> Iterable[openai.types.chat.completion_create_params.Function]:
+def capabilities_to_functions(
+    capabilities: Dict[str, Capability],
+) -> Iterable[openai.types.chat.completion_create_params.Function]:
     """
     This function takes a dictionary of capabilities and returns a dictionary of functions, that can be called with the
     parameters of the respective capabilities.
@@ -186,13 +204,21 @@ def capabilities_to_functions(capabilities: Dict[str, Capability]) -> Iterable[o
     ]
 
 
-def capabilities_to_tools(capabilities: Dict[str, Capability]) -> Iterable[openai.types.chat.completion_create_params.ChatCompletionToolParam]:
+def capabilities_to_tools(
+    capabilities: Dict[str, Capability],
+) -> Iterable[openai.types.chat.completion_create_params.ChatCompletionToolParam]:
     """
     This function takes a dictionary of capabilities and returns a dictionary of functions, that can be called with the
     parameters of the respective capabilities.
     """
     return [
-        ChatCompletionToolParam(type="function", function=Function(name=name, description=capability.describe(), parameters=capability.to_model().model_json_schema()))
+        ChatCompletionToolParam(
+            type="function",
+            function=Function(
+                name=name,
+                description=capability.describe(),
+                parameters=capability.to_model().model_json_schema(),
+            ),
+        )
         for name, capability in capabilities.items()
     ]
-