Merge branch 'development' into merge_web_api_testing_development

Author: DianaStrauss

README.md

@@ -17,7 +17,7 @@ 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.4.0-dev"
+version = "0.4.0"
 license = { file = "LICENSE" }
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -28,20 +28,20 @@ classifiers = [
 dependencies = [
     'fabric == 3.2.2',
     'Mako == 1.3.2',
-    'requests == 2.32.0',
+    'requests == 2.32.3',
     'rich == 13.7.1',
     'tiktoken == 0.8.0',
-    'instructor == 1.3.5',
+    'instructor == 1.7.2',
     'PyYAML == 6.0.1',
     'python-dotenv == 1.0.1',
     'pypsexec == 0.3.0',
     'pydantic == 2.8.2',
-    'openai == 1.28.0',
+    'openai == 1.65.2',
     'BeautifulSoup4',
     'nltk',
     'fastapi == 0.114.0',
     'fastapi-utils == 0.7.0',
-    'jinja2 == 3.1.4',
+    'jinja2 == 3.1.6',
     'uvicorn[standard] == 0.30.6',
     'dataclasses_json == 0.6.7',
     'websockets == 13.1',
@@ -52,7 +52,9 @@ dependencies = [
     'langchain_community',
     'langchain_chroma',
     'langchain_openai'
-
+    'langchain-openai',
+    'markdown',
+    'chromadb',
 ]
 
 [project.urls]

pyproject.toml

@@ -2,19 +2,22 @@
 import sys
 
 from hackingBuddyGPT.usecases.base import use_cases
+from hackingBuddyGPT.utils.configurable import CommandMap, InvalidCommand, Parseable, instantiate
 
 
 def main():
-    parser = argparse.ArgumentParser()
-    subparser = parser.add_subparsers(required=True)
-    for name, use_case in use_cases.items():
-        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(configuration=configuration)
-    instance.run()
+    use_case_parsers: CommandMap = {
+        name: Parseable(use_case, description=use_case.description)
+        for name, use_case in use_cases.items()
+    }
+    try:
+        instance, configuration = instantiate(sys.argv, use_case_parsers)
+    except InvalidCommand as e:
+        if len(f"{e}") > 0:
+            print(e)
+        print(e.usage)
+        sys.exit(1)
+    instance.run(configuration)
 
 
 if __name__ == "__main__":

src/hackingBuddyGPT/cli/wintermute.py

@@ -4,7 +4,7 @@
 from mako.template import Template
 from typing import Dict
 
-from hackingBuddyGPT.utils.logging import log_conversation, GlobalLogger
+from hackingBuddyGPT.utils.logging import log_conversation, Logger, log_param
 from hackingBuddyGPT.capabilities.capability import (
     Capability,
     capabilities_to_simple_text_handler,
@@ -15,7 +15,7 @@
 
 @dataclass
 class Agent(ABC):
-    log: GlobalLogger = None
+    log: Logger = log_param
 
     _capabilities: Dict[str, Capability] = field(default_factory=dict)
     _default_capability: Capability = None

src/hackingBuddyGPT/usecases/agents.py

@@ -3,12 +3,10 @@
 import argparse
 from dataclasses import dataclass
 
-from hackingBuddyGPT.utils.logging import GlobalLogger
+from hackingBuddyGPT.utils.logging import Logger, log_param
 from typing import Dict, Type, TypeVar, Generic
 
-from hackingBuddyGPT.utils.configurable import ParameterDefinitions, build_parser, get_arguments, get_class_parameters, \
-    Transparent, ParserState
-
+from hackingBuddyGPT.utils.configurable import Transparent, configurable
 
 @dataclass
 class UseCase(abc.ABC):
@@ -22,22 +20,21 @@ class UseCase(abc.ABC):
     so that they can be automatically discovered and run from the command line.
     """
 
-    log: GlobalLogger
+    log: Logger = log_param
 
-    def init(self, configuration):
+    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.
         """
-        self.configuration = configuration
-        self.log.start_run(self.get_name(), self.serialize_configuration(configuration))
+        pass
 
     def serialize_configuration(self, configuration) -> str:
         return json.dumps(configuration)
 
     @abc.abstractmethod
-    def run(self):
+    def run(self, configuration):
         """
         The run method is the main method of the UseCase. It is used to run the UseCase, and should contain the main
         logic. It is recommended to have only the main llm loop in here, and call out to other methods for the
@@ -70,7 +67,10 @@ def before_run(self):
     def after_run(self):
         pass
 
-    def run(self):
+    def run(self, configuration):
+        self.configuration = configuration
+        self.log.start_run(self.get_name(), self.serialize_configuration(configuration))
+
         self.before_run()
 
         turn = 1
@@ -98,31 +98,10 @@ def run(self):
             raise
 
 
-@dataclass
-class _WrappedUseCase:
-    """
-    A WrappedUseCase should not be used directly and is an internal tool used for initialization and dependency injection
-    of the actual UseCases.
-    """
-
-    name: str
-    description: str
-    use_case: Type[UseCase]
-    parameters: ParameterDefinitions
-
-    def build_parser(self, parser: argparse.ArgumentParser):
-        parser_state = ParserState()
-        build_parser(self.parameters, parser, parser_state)
-        parser.set_defaults(use_case=self, parser_state=parser_state)
-
-    def __call__(self, args: argparse.Namespace):
-        return self.use_case(**get_arguments(self.parameters, args, args.parser_state))
-
-
-use_cases: Dict[str, _WrappedUseCase] = dict()
+use_cases: Dict[str, configurable] = dict()
 
 
-T = TypeVar("T")
+T = TypeVar("T", bound=type)
 
 
 class AutonomousAgentUseCase(AutonomousUseCase, Generic[T]):
@@ -137,13 +116,12 @@ def get_name(self) -> str:
     @classmethod
     def __class_getitem__(cls, item):
         item = dataclass(item)
-        item.__parameters__ = get_class_parameters(item)
 
         class AutonomousAgentUseCase(AutonomousUseCase):
             agent: Transparent(item) = None
 
-            def init(self, configuration):
-                super().init(configuration)
+            def init(self):
+                super().init()
                 self.agent.init()
 
             def get_name(self) -> str:
@@ -169,7 +147,7 @@ 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, description, cls, get_class_parameters(cls))
+        use_cases[name] = configurable(name, description)(cls)
         return cls
 
     return inner
@@ -181,4 +159,4 @@ def register_use_case(name: str, description: str, use_case: Type[UseCase]):
     """
     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))
+    use_cases[name] = configurable(name, description)(use_case)

src/hackingBuddyGPT/usecases/base.py

@@ -8,8 +8,8 @@
 class ExPrivEscLinuxHintFileUseCase(AutonomousAgentUseCase[LinuxPrivesc]):
     hints: str = None
 
-    def init(self, configuration):
-        super().init(configuration)
+    def init(self):
+        super().init()
         self.agent.hint = self.read_hint()
 
     # simple helper that reads the hints file and returns the hint

src/hackingBuddyGPT/usecases/examples/hintfile.py

@@ -2,10 +2,10 @@
 This usecase is an extension of `usecase/privesc`.
 
 ## Setup
-### Depdendencies
+### Dependencies
 The needed dependencies can be downloaded with `pip install -e '.[rag-usecase]'`. If you encounter the error `unexpected keyword argument 'proxies'` after trying to start the usecase, try downgrading `httpx` to 0.27.2.
 ### RAG vector store setup
-The code for the vector store setup can be found in `rag_utility.py`. Currently the vectore store uses two sources: `GTFObins` and `hacktricks`. To use RAG, download the markdown files and place them in `rag_storage/GTFObinMarkdownfiles` (`rag_storage/hacktricksMarkdownFiles`). You can download the markdown files either from the respective github repository ([GTFObin](https://github.com/GTFOBins/GTFOBins.github.io/tree/master), [hacktricks](https://github.com/HackTricks-wiki/hacktricks/tree/master/src/linux-hardening/privilege-escalation)) or scrape them from their website ([GTFObin](https://gtfobins.github.io/), [hacktricks](https://book.hacktricks.wiki/en/linux-hardening/privilege-escalation/index.html)).
+The code for the vector store setup can be found in `rag_utility.py`. Currently the vector store uses two sources: `GTFObins` and `hacktricks`. To use RAG, download the markdown files and place them in `rag_storage/GTFObinMarkdownfiles` (`rag_storage/hacktricksMarkdownFiles`). You can download the markdown files either from the respective github repository ([GTFObin](https://github.com/GTFOBins/GTFOBins.github.io/tree/master), [hacktricks](https://github.com/HackTricks-wiki/hacktricks/tree/master/src/linux-hardening/privilege-escalation)) or scrape them from their website ([GTFObin](https://gtfobins.github.io/), [hacktricks](https://book.hacktricks.wiki/en/linux-hardening/privilege-escalation/index.html)).
 
 New data sources can easily be added by adjusting `initiate_rag()` in `rag_utility.py`.
 

src/hackingBuddyGPT/usecases/rag/README.md

@@ -20,8 +20,8 @@ def init(self):
 class ThesisLinuxPrivescPrototypeUseCase(AutonomousAgentUseCase[ThesisLinuxPrivescPrototype]):
     hints: str = ""
 
-    def init(self,configuration):
-        super().init(configuration)
+    def init(self):
+        super().init()
         if self.hints != "":
             self.agent.hint = self.read_hint()
 

src/hackingBuddyGPT/usecases/rag/linux.py

@@ -18,6 +18,7 @@
 from starlette.templating import Jinja2Templates
 
 from hackingBuddyGPT.usecases.base import UseCase, use_case
+from hackingBuddyGPT.utils.configurable import parameter
 from hackingBuddyGPT.utils.db_storage import DbStorage
 from hackingBuddyGPT.utils.db_storage.db_storage import (
     Message,
@@ -205,10 +206,9 @@ class Viewer(UseCase):
     TODOs:
     - [ ] This server needs to be as async as possible to allow good performance, but the database accesses are not yet, might be an issue?
     """
-    log: GlobalLocalLogger
-    log_db: DbStorage
-    listen_host: str = "127.0.0.1"
-    listen_port: int = 4444
+    log: GlobalLocalLogger = None
+    log_db: DbStorage = None
+    log_server_address: str = "127.0.0.1:4444"
     save_playback_dir: str = ""
 
     async def save_message(self, message: ControlMessage):
@@ -232,7 +232,7 @@ async def save_message(self, message: ControlMessage):
         with open(file_path, "a") as f:
             f.write(ReplayMessage(datetime.datetime.now(), message).to_json() + "\n")
 
-    def run(self):
+    def run(self, config):
         @asynccontextmanager
         async def lifespan(app: FastAPI):
             app.state.db = self.log_db
@@ -337,16 +337,30 @@ async def client_endpoint(websocket: WebSocket):
                 print("Egress WebSocket disconnected")
 
         import uvicorn
-        uvicorn.run(app, host=self.listen_host, port=self.listen_port)
+        listen_parts = self.log_server_address.split(":", 1)
+        if len(listen_parts) != 2:
+            if listen_parts[0].startswith("http://"):
+                listen_parts.append("80")
+            elif listen_parts[0].startswith("https://"):
+                listen_parts.append("443")
+            else:
+                raise ValueError(f"Invalid log server address (does not contain http/https or a port): {self.log_server_address}")
+
+        listen_host, listen_port = listen_parts[0], int(listen_parts[1])
+        if listen_host.startswith("http://"):
+            listen_host = listen_host[len("http://"):]
+        elif listen_host.startswith("https://"):
+            listen_host = listen_host[len("https://"):]
+        uvicorn.run(app, host=listen_host, port=listen_port)
 
     def get_name(self) -> str:
         return "log_viewer"
 
 
 @use_case("Tool to replay the .jsonl logs generated by the Viewer (not well tested)")
 class Replayer(UseCase):
-    log: GlobalRemoteLogger
-    replay_file: str
+    log: GlobalRemoteLogger = None
+    replay_file: str = None
     pause_on_message: bool = False
     pause_on_tool_calls: bool = False
     playback_speed: float = 1.0

src/hackingBuddyGPT/usecases/viewer.py

@@ -1,4 +1,4 @@
-from .configurable import Configurable, configurable
+from .configurable import Configurable, configurable, parameter
 from .console import *
 from .db_storage import *
 from .llm_util import *