Restructure

.gitignore

@@ -7,3 +7,6 @@ __pycache__/
 *.sqlite3
 *.sqlite3-jounal
 *.sqlite
+src/hackingBuddyGPT.egg-info/
+build/
+dist/

pyproject.toml

@@ -0,0 +1,58 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "hackingBuddyGPT"
+authors = [
+	{ name = "Andreas Happe", email = "[email protected]" }
+]
+maintainers = [
+	{ name = "Andreas Happe", email = "[email protected]" },
+	{ name = "Juergen Cito", email = "[email protected]" }
+]
+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"
+version = "0.3.0"
+license = { file = "LICENSE" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+	"Development Status :: 4 - Beta",
+]
+dependencies = [
+	'fabric == 3.2.2',
+	'Mako == 1.3.2',
+	'requests == 2.32.0',
+	'rich == 13.7.1',
+	'tiktoken == 0.6.0',
+	'instructor == 1.2.2',
+	'PyYAML == 6.0.1',
+	'python-dotenv == 1.0.1',
+	'pypsexec == 0.3.0'
+]
+
+[project.urls]
+Homepage = "https://www.hackingbuddy.ai"
+Documentation = "https://docs.hackingbuddy.ai"
+Repository = "https://github.com/ipa-lab/hackingBuddyGPT"
+"Bug Tracker" = "https://github.com/ipa-lab/hackingBuddyGPT/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+"hackingBuddyGPT.usecases.privesc.templates" = ["*.txt"]
+"hackingBuddyGPT.usecases.minimal" = ["*.txt"]
+
+[tool.pytest.ini_options]
+pythonpath = "src"
+addopts = [
+	"--import-mode=importlib",
+]
+
+[project.scripts]
+wintermute = "hackingBuddyGPT.cli.wintermute:main"

capabilities/http_request.py → src/hackingBuddyGPT/capabilities/http_request.py

@@ -1,11 +1,9 @@
 import base64
 from dataclasses import dataclass
-from typing import Literal, Optional, Dict
-
 import requests
+from typing import Literal, Optional, Dict
 
-from capabilities import Capability
-
+from . import Capability
 
 @dataclass
 class HTTPRequest(Capability):

capabilities/psexec_run_command.py → src/hackingBuddyGPT/capabilities/psexec_run_command.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Tuple
 
-from utils import PSExecConnection
+from hackingBuddyGPT.utils import PSExecConnection
 from .capability import Capability
 
 

capabilities/psexec_test_credential.py → src/hackingBuddyGPT/capabilities/psexec_test_credential.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Tuple
 
-from utils import PSExecConnection
+from hackingBuddyGPT.utils import PSExecConnection
 from .capability import Capability
 
 

capabilities/record_note.py → src/hackingBuddyGPT/capabilities/record_note.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Tuple, List
 
-from capabilities import Capability
+from . import Capability
 
 
 @dataclass

capabilities/ssh_run_command.py → src/hackingBuddyGPT/capabilities/ssh_run_command.py

@@ -1,19 +1,13 @@
 import re
-from dataclasses import dataclass
-from typing import Tuple
 
+from dataclasses import dataclass
 from invoke import Responder
-
 from io import StringIO
-from utils import SSHConnection
-from .capability import Capability
-
-
-GOT_ROOT_REXEXPs = [
-    re.compile("^# $"),
-    re.compile("^bash-[0-9]+.[0-9]# $")
-]
+from typing import Tuple
 
+from hackingBuddyGPT.utils import SSHConnection
+from hackingBuddyGPT.utils.shell_root_detection import got_root
+from .capability import Capability
 
 @dataclass
 class SSHRunCommand(Capability):
@@ -27,8 +21,6 @@ def get_name(self):
         return "exec_command"
 
     def __call__(self, command: str) -> Tuple[str, bool]:
-        got_root = False
-
         if command.startswith(self.get_name()):
             cmd_parts = command.split(" ", 1)
             command = cmd_parts[1]
@@ -57,9 +49,4 @@ def __call__(self, command: str) -> Tuple[str, bool]:
         ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
         last_line = ansi_escape.sub('', last_line)
 
-        for i in GOT_ROOT_REXEXPs:
-            if i.fullmatch(last_line):
-                got_root = True
-        if last_line.startswith(f'root@{self.conn.hostname}:'):
-            got_root = True
-        return tmp, got_root
+        return tmp, got_root(self.conn.hostname, last_line)

capabilities/ssh_test_credential.py → src/hackingBuddyGPT/capabilities/ssh_test_credential.py

@@ -3,7 +3,7 @@
 
 import paramiko
 
-from utils import SSHConnection
+from hackingBuddyGPT.utils import SSHConnection
 from .capability import Capability
 
 

capabilities/submit_flag.py → src/hackingBuddyGPT/capabilities/submit_flag.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Tuple, List, Set, Callable
 
-from capabilities import Capability
+from . import Capability
 
 
 @dataclass

wintermute.py → src/hackingBuddyGPT/cli/wintermute.py

@@ -1,7 +1,7 @@
 import argparse
 import sys
 
-from usecases.base import use_cases
+from hackingBuddyGPT.usecases.base import use_cases
 
 
 def main():

usecases/agents.py → src/hackingBuddyGPT/usecases/agents.py

@@ -1,13 +1,13 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
+from mako.template import Template
+from rich.panel import Panel
 from typing import Dict
 
-from capabilities.capability import Capability, capabilities_to_simple_text_handler
-from usecases.common_patterns import RoundBasedUseCase
+from hackingBuddyGPT.utils import llm_util
 
-from mako.template import Template
-from rich.panel import Panel
-from utils import llm_util
+from hackingBuddyGPT.capabilities.capability import Capability, capabilities_to_simple_text_handler
+from .common_patterns import RoundBasedUseCase
 
 @dataclass
 class Agent(RoundBasedUseCase, ABC):
@@ -50,7 +50,6 @@ def init(self):
         super().init()
 
     def set_initial_state(self, initial_state):
-        print("setting state!")
         self._state = initial_state
 
     def set_template(self, template):

usecases/base.py → src/hackingBuddyGPT/usecases/base.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass, field
 from typing import Dict, Type
 
-from utils.configurable import ParameterDefinitions, build_parser, get_arguments, get_class_parameters
+from hackingBuddyGPT.utils.configurable import ParameterDefinitions, build_parser, get_arguments, get_class_parameters
 
 class UseCase(abc.ABC):
     """

usecases/common_patterns.py → src/hackingBuddyGPT/usecases/common_patterns.py

@@ -2,9 +2,10 @@
 
 from dataclasses import dataclass
 from rich.panel import Panel
-from usecases.base import UseCase
-from utils import Console, DbStorage
-from utils.openai.openai_llm import OpenAIConnection
+
+from .base import UseCase
+from hackingBuddyGPT.utils import Console, DbStorage
+from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
 
 # this set ups all the console and database stuff, and runs the main loop for a bounded amount of turns
 @dataclass

usecases/minimal/agent.py → src/hackingBuddyGPT/usecases/minimal/agent.py

@@ -1,14 +1,13 @@
 import pathlib
 from dataclasses import dataclass, field
-
 from mako.template import Template
 from rich.panel import Panel
 
-from capabilities import SSHRunCommand, SSHTestCredential
-from utils import SSHConnection, llm_util
-from usecases.base import use_case
-from usecases.agents import Agent
-from utils.cli_history import SlidingCliHistory
+from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential
+from hackingBuddyGPT.utils import SSHConnection, llm_util
+from hackingBuddyGPT.usecases.base import use_case
+from hackingBuddyGPT.usecases.agents import Agent
+from hackingBuddyGPT.utils.cli_history import SlidingCliHistory
 
 template_dir = pathlib.Path(__file__).parent
 template_next_cmd = Template(filename=str(template_dir / "next_cmd.txt"))

usecases/minimal/agent_with_state.py → src/hackingBuddyGPT/usecases/minimal/agent_with_state.py

@@ -2,11 +2,11 @@
 import pathlib
 from dataclasses import dataclass
 
-from capabilities import SSHRunCommand, SSHTestCredential
-from utils import SSHConnection, llm_util
-from usecases.base import use_case
-from usecases.agents import TemplatedAgent, AgentWorldview
-from utils.cli_history import SlidingCliHistory
+from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential
+from hackingBuddyGPT.utils import SSHConnection, llm_util
+from hackingBuddyGPT.usecases.base import use_case
+from hackingBuddyGPT.usecases.agents import TemplatedAgent, AgentWorldview
+from hackingBuddyGPT.utils.cli_history import SlidingCliHistory
 
 @dataclass
 class MinimalLinuxTemplatedPrivescState(AgentWorldview):
@@ -47,4 +47,4 @@ def init(self):
 
         # setup state
         max_history_size = self.llm.context_size - llm_util.SAFETY_MARGIN - self._template_size
-        self.set_initial_state(MinimalLinuxTemplatedPrivescState(self.conn, self.llm, max_history_size))
+        self.set_initial_state(MinimalLinuxTemplatedPrivescState(self.conn, self.llm, max_history_size))

usecases/privesc/common.py → src/hackingBuddyGPT/usecases/privesc/common.py

@@ -1,15 +1,14 @@
 import pathlib
 from dataclasses import dataclass, field
-from typing import Dict
-
 from mako.template import Template
 from rich.panel import Panel
+from typing import Dict
 
-from capabilities import Capability
-from capabilities.capability import capabilities_to_simple_text_handler
-from usecases.agents import Agent
-from utils import llm_util, ui
-from utils.cli_history import SlidingCliHistory
+from hackingBuddyGPT.capabilities import Capability
+from hackingBuddyGPT.capabilities.capability import capabilities_to_simple_text_handler
+from hackingBuddyGPT.usecases.agents import Agent
+from hackingBuddyGPT.utils import llm_util, ui
+from hackingBuddyGPT.utils.cli_history import SlidingCliHistory
 
 template_dir = pathlib.Path(__file__).parent / "templates"
 template_next_cmd = Template(filename=str(template_dir / "query_next_command.txt"))

usecases/privesc/linux.py → src/hackingBuddyGPT/usecases/privesc/linux.py

@@ -1,16 +1,15 @@
 import json
 import pathlib
 from dataclasses import dataclass
-
 from mako.template import Template
 
-from capabilities import SSHRunCommand, SSHTestCredential
-from usecases.privesc.common import Privesc
-from utils import SSHConnection
-from usecases.base import use_case, UseCase
-from utils.console.console import Console
-from utils.db_storage.db_storage import DbStorage
-from utils.openai.openai_llm import OpenAIConnection
+from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential
+from .common import Privesc
+from hackingBuddyGPT.utils import SSHConnection
+from hackingBuddyGPT.usecases.base import use_case, UseCase
+from hackingBuddyGPT.utils.console.console import Console
+from hackingBuddyGPT.utils.db_storage.db_storage import DbStorage
+from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection
 
 template_dir = pathlib.Path(__file__).parent / "templates"
 template_next_cmd = Template(filename=str(template_dir / "query_next_command.txt"))

usecases/privesc/windows.py → src/hackingBuddyGPT/usecases/privesc/windows.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
-from capabilities.psexec_run_command import PSExecRunCommand
-from capabilities.psexec_test_credential import PSExecTestCredential
-from usecases.base import use_case
-from usecases.privesc.common import Privesc
-from utils.psexec.psexec import PSExecConnection
+
+from hackingBuddyGPT.capabilities.psexec_run_command import PSExecRunCommand
+from hackingBuddyGPT.capabilities.psexec_test_credential import PSExecTestCredential
+from hackingBuddyGPT.usecases.base import use_case
+from hackingBuddyGPT.usecases.privesc.common import Privesc
+from hackingBuddyGPT.utils.psexec.psexec import PSExecConnection
 
 
 @use_case("windows_privesc", "Windows Privilege Escalation")