diff --git a/.gitignore b/.gitignore index 9b7eff2c..44a3be8e 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,6 @@ __pycache__/ *.sqlite3 *.sqlite3-jounal *.sqlite +src/hackingBuddyGPT.egg-info/ +build/ +dist/ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..ecb1d309 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,58 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "hackingBuddyGPT" +authors = [ + { name = "Andreas Happe", email = "andreas@offensive.one" } +] +maintainers = [ + { name = "Andreas Happe", email = "andreas@offensive.one" }, + { name = "Juergen Cito", email = "juergen.cito@tuwiena.c.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.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" diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index c00660db..00000000 --- a/requirements.txt +++ /dev/null @@ -1,30 +0,0 @@ -bcrypt==4.1.2 -certifi==2024.2.2 -cffi==1.16.0 -charset-normalizer==3.3.2 -cryptography==42.0.5 -decorator==5.1.1 -Deprecated==1.2.14 -fabric==3.2.2 -idna==3.7 -invoke==2.2.0 -Mako==1.3.2 -markdown-it-py==3.0.0 -MarkupSafe==2.1.5 -mdurl==0.1.2 -paramiko==3.4.0 -pycparser==2.22 -Pygments==2.17.2 -PyNaCl==1.5.0 -pypsexec==0.3.0 -pyspnego==0.10.2 -python-dotenv==1.0.1 -regex==2023.12.25 -requests==2.32.0 -rich==13.7.1 -smbprotocol==1.13.0 -tiktoken==0.6.0 -urllib3==2.2.1 -wrapt==1.16.0 -instructor==1.2.2 -PyYAML==6.0.1 diff --git a/src/hackingBuddyGPT/__init__.py b/src/hackingBuddyGPT/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/capabilities/__init__.py b/src/hackingBuddyGPT/capabilities/__init__.py similarity index 100% rename from capabilities/__init__.py rename to src/hackingBuddyGPT/capabilities/__init__.py diff --git a/capabilities/capability.py b/src/hackingBuddyGPT/capabilities/capability.py similarity index 100% rename from capabilities/capability.py rename to src/hackingBuddyGPT/capabilities/capability.py diff --git a/capabilities/http_request.py b/src/hackingBuddyGPT/capabilities/http_request.py similarity index 97% rename from capabilities/http_request.py rename to src/hackingBuddyGPT/capabilities/http_request.py index 2088124e..533eae10 100644 --- a/capabilities/http_request.py +++ b/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): diff --git a/capabilities/psexec_run_command.py b/src/hackingBuddyGPT/capabilities/psexec_run_command.py similarity index 92% rename from capabilities/psexec_run_command.py rename to src/hackingBuddyGPT/capabilities/psexec_run_command.py index 85456e52..f0a47913 100644 --- a/capabilities/psexec_run_command.py +++ b/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 diff --git a/capabilities/psexec_test_credential.py b/src/hackingBuddyGPT/capabilities/psexec_test_credential.py similarity index 94% rename from capabilities/psexec_test_credential.py rename to src/hackingBuddyGPT/capabilities/psexec_test_credential.py index 173ec683..7cebcaaf 100644 --- a/capabilities/psexec_test_credential.py +++ b/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 diff --git a/capabilities/record_note.py b/src/hackingBuddyGPT/capabilities/record_note.py similarity index 92% rename from capabilities/record_note.py rename to src/hackingBuddyGPT/capabilities/record_note.py index 3e69614e..7e773125 100644 --- a/capabilities/record_note.py +++ b/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 diff --git a/capabilities/ssh_run_command.py b/src/hackingBuddyGPT/capabilities/ssh_run_command.py similarity index 80% rename from capabilities/ssh_run_command.py rename to src/hackingBuddyGPT/capabilities/ssh_run_command.py index 4721856d..c0a30ff0 100644 --- a/capabilities/ssh_run_command.py +++ b/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) diff --git a/capabilities/ssh_test_credential.py b/src/hackingBuddyGPT/capabilities/ssh_test_credential.py similarity index 94% rename from capabilities/ssh_test_credential.py rename to src/hackingBuddyGPT/capabilities/ssh_test_credential.py index e64814fd..2f6dd4bb 100644 --- a/capabilities/ssh_test_credential.py +++ b/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 diff --git a/capabilities/submit_flag.py b/src/hackingBuddyGPT/capabilities/submit_flag.py similarity index 97% rename from capabilities/submit_flag.py rename to src/hackingBuddyGPT/capabilities/submit_flag.py index 4f48b259..b481fd97 100644 --- a/capabilities/submit_flag.py +++ b/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 diff --git a/src/hackingBuddyGPT/cli/__init__.py b/src/hackingBuddyGPT/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/stats.py b/src/hackingBuddyGPT/cli/stats.py similarity index 100% rename from stats.py rename to src/hackingBuddyGPT/cli/stats.py diff --git a/viewer.py b/src/hackingBuddyGPT/cli/viewer.py similarity index 100% rename from viewer.py rename to src/hackingBuddyGPT/cli/viewer.py diff --git a/wintermute.py b/src/hackingBuddyGPT/cli/wintermute.py similarity index 90% rename from wintermute.py rename to src/hackingBuddyGPT/cli/wintermute.py index b0e3cb62..85552b3b 100644 --- a/wintermute.py +++ b/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(): diff --git a/usecases/__init__.py b/src/hackingBuddyGPT/usecases/__init__.py similarity index 100% rename from usecases/__init__.py rename to src/hackingBuddyGPT/usecases/__init__.py diff --git a/usecases/agents.py b/src/hackingBuddyGPT/usecases/agents.py similarity index 93% rename from usecases/agents.py rename to src/hackingBuddyGPT/usecases/agents.py index ae3fb5ec..7c5f1f1a 100644 --- a/usecases/agents.py +++ b/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): diff --git a/usecases/base.py b/src/hackingBuddyGPT/usecases/base.py similarity index 95% rename from usecases/base.py rename to src/hackingBuddyGPT/usecases/base.py index 2775e01b..f090e4e8 100644 --- a/usecases/base.py +++ b/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): """ diff --git a/usecases/common_patterns.py b/src/hackingBuddyGPT/usecases/common_patterns.py similarity index 91% rename from usecases/common_patterns.py rename to src/hackingBuddyGPT/usecases/common_patterns.py index df273b1e..357a56f7 100644 --- a/usecases/common_patterns.py +++ b/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 diff --git a/usecases/minimal/__init__.py b/src/hackingBuddyGPT/usecases/minimal/__init__.py similarity index 100% rename from usecases/minimal/__init__.py rename to src/hackingBuddyGPT/usecases/minimal/__init__.py diff --git a/usecases/minimal/agent.py b/src/hackingBuddyGPT/usecases/minimal/agent.py similarity index 86% rename from usecases/minimal/agent.py rename to src/hackingBuddyGPT/usecases/minimal/agent.py index e8e5034f..555a0684 100644 --- a/usecases/minimal/agent.py +++ b/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")) diff --git a/usecases/minimal/agent_with_state.py b/src/hackingBuddyGPT/usecases/minimal/agent_with_state.py similarity index 79% rename from usecases/minimal/agent_with_state.py rename to src/hackingBuddyGPT/usecases/minimal/agent_with_state.py index d75cbd12..35a1f217 100644 --- a/usecases/minimal/agent_with_state.py +++ b/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)) \ No newline at end of file + self.set_initial_state(MinimalLinuxTemplatedPrivescState(self.conn, self.llm, max_history_size)) diff --git a/usecases/minimal/next_cmd.txt b/src/hackingBuddyGPT/usecases/minimal/next_cmd.txt similarity index 100% rename from usecases/minimal/next_cmd.txt rename to src/hackingBuddyGPT/usecases/minimal/next_cmd.txt diff --git a/usecases/privesc/__init__.py b/src/hackingBuddyGPT/usecases/privesc/__init__.py similarity index 100% rename from usecases/privesc/__init__.py rename to src/hackingBuddyGPT/usecases/privesc/__init__.py diff --git a/usecases/privesc/common.py b/src/hackingBuddyGPT/usecases/privesc/common.py similarity index 94% rename from usecases/privesc/common.py rename to src/hackingBuddyGPT/usecases/privesc/common.py index 340837f5..9778682b 100644 --- a/usecases/privesc/common.py +++ b/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")) diff --git a/usecases/privesc/linux.py b/src/hackingBuddyGPT/usecases/privesc/linux.py similarity index 92% rename from usecases/privesc/linux.py rename to src/hackingBuddyGPT/usecases/privesc/linux.py index 4b4a88d5..ccd1065f 100644 --- a/usecases/privesc/linux.py +++ b/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")) diff --git a/usecases/privesc/templates/analyze_cmd.txt b/src/hackingBuddyGPT/usecases/privesc/templates/analyze_cmd.txt similarity index 100% rename from usecases/privesc/templates/analyze_cmd.txt rename to src/hackingBuddyGPT/usecases/privesc/templates/analyze_cmd.txt diff --git a/usecases/privesc/templates/get_hint_from_lse.txt b/src/hackingBuddyGPT/usecases/privesc/templates/get_hint_from_lse.txt similarity index 100% rename from usecases/privesc/templates/get_hint_from_lse.txt rename to src/hackingBuddyGPT/usecases/privesc/templates/get_hint_from_lse.txt diff --git a/usecases/privesc/templates/query_next_command.txt b/src/hackingBuddyGPT/usecases/privesc/templates/query_next_command.txt similarity index 100% rename from usecases/privesc/templates/query_next_command.txt rename to src/hackingBuddyGPT/usecases/privesc/templates/query_next_command.txt diff --git a/usecases/privesc/templates/update_state.txt b/src/hackingBuddyGPT/usecases/privesc/templates/update_state.txt similarity index 100% rename from usecases/privesc/templates/update_state.txt rename to src/hackingBuddyGPT/usecases/privesc/templates/update_state.txt diff --git a/usecases/privesc/windows.py b/src/hackingBuddyGPT/usecases/privesc/windows.py similarity index 53% rename from usecases/privesc/windows.py rename to src/hackingBuddyGPT/usecases/privesc/windows.py index 463e41a6..cf565095 100644 --- a/usecases/privesc/windows.py +++ b/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") diff --git a/usecases/web/__init__.py b/src/hackingBuddyGPT/usecases/web/__init__.py similarity index 100% rename from usecases/web/__init__.py rename to src/hackingBuddyGPT/usecases/web/__init__.py diff --git a/usecases/web/simple.py b/src/hackingBuddyGPT/usecases/web/simple.py similarity index 86% rename from usecases/web/simple.py rename to src/hackingBuddyGPT/usecases/web/simple.py index c44db107..c9177d3f 100644 --- a/usecases/web/simple.py +++ b/src/hackingBuddyGPT/usecases/web/simple.py @@ -1,21 +1,21 @@ +import pydantic_core import time -from dataclasses import dataclass, field -from typing import List, Any, Union, Dict -import pydantic_core +from dataclasses import dataclass, field from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage from rich.panel import Panel +from typing import List, Any, Union, Dict -from capabilities import Capability -from capabilities.capability import capabilities_to_action_model -from capabilities.http_request import HTTPRequest -from capabilities.record_note import RecordNote -from capabilities.submit_flag import SubmitFlag -from utils import LLMResult, tool_message -from usecases.base import use_case -from usecases.common_patterns import RoundBasedUseCase -from utils.configurable import parameter -from utils.openai.openai_lib import OpenAILib +from hackingBuddyGPT.capabilities import Capability +from hackingBuddyGPT.capabilities.capability import capabilities_to_action_model +from hackingBuddyGPT.capabilities.http_request import HTTPRequest +from hackingBuddyGPT.capabilities.record_note import RecordNote +from hackingBuddyGPT.capabilities.submit_flag import SubmitFlag +from hackingBuddyGPT.utils import LLMResult, tool_message +from hackingBuddyGPT.usecases.base import use_case +from hackingBuddyGPT.usecases.common_patterns import RoundBasedUseCase +from hackingBuddyGPT.utils.configurable import parameter +from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib Prompt = List[Union[ChatCompletionMessage, ChatCompletionMessageParam]] diff --git a/usecases/web_api_testing/__init__.py b/src/hackingBuddyGPT/usecases/web_api_testing/__init__.py similarity index 100% rename from usecases/web_api_testing/__init__.py rename to src/hackingBuddyGPT/usecases/web_api_testing/__init__.py diff --git a/usecases/web_api_testing/prompt_engineer.py b/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py similarity index 99% rename from usecases/web_api_testing/prompt_engineer.py rename to src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py index 6648c2f7..5d7fcf8b 100644 --- a/usecases/web_api_testing/prompt_engineer.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py @@ -1,6 +1,6 @@ from openai.types.chat import ChatCompletionMessage -from utils import openai +from hackingBuddyGPT.utils import openai class PromptEngineer(object): '''Prompt engineer that creates prompts of different types''' diff --git a/usecases/web_api_testing/simple_openapi_documentation.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py similarity index 93% rename from usecases/web_api_testing/simple_openapi_documentation.py rename to src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py index 5ec0a90f..03b34cbd 100644 --- a/usecases/web_api_testing/simple_openapi_documentation.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py @@ -1,30 +1,29 @@ import datetime import os +import pydantic_core import time -from dataclasses import dataclass, field -from typing import List, Any, Union, Dict +import yaml +from dataclasses import dataclass, field from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage -from capabilities import Capability -from capabilities.capability import capabilities_to_action_model -from capabilities.http_request import HTTPRequest -from capabilities.record_note import RecordNote -from capabilities.submit_flag import SubmitFlag -from usecases.common_patterns import RoundBasedUseCase -from usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy -from utils import LLMResult, tool_message, ui -from utils.configurable import parameter -from utils.openai.openai_lib import OpenAILib from rich.panel import Panel -from usecases import use_case +from typing import List, Any, Union, Dict -import pydantic_core -import yaml +from hackingBuddyGPT.capabilities import Capability +from hackingBuddyGPT.capabilities.capability import capabilities_to_action_model +from hackingBuddyGPT.capabilities.http_request import HTTPRequest +from hackingBuddyGPT.capabilities.record_note import RecordNote +from hackingBuddyGPT.capabilities.submit_flag import SubmitFlag +from hackingBuddyGPT.usecases.common_patterns import RoundBasedUseCase +from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy +from hackingBuddyGPT.utils import LLMResult, tool_message, ui +from hackingBuddyGPT.utils.configurable import parameter +from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib +from hackingBuddyGPT.usecases import use_case Prompt = List[Union[ChatCompletionMessage, ChatCompletionMessageParam]] Context = Any - @use_case("simple_web_api_documentation", "Minimal implementation of a web api documentation use case") @dataclass class SimpleWebAPIDocumentation(RoundBasedUseCase): diff --git a/usecases/web_api_testing/simple_web_api_testing.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py similarity index 89% rename from usecases/web_api_testing/simple_web_api_testing.py rename to src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py index 1b3eaa8a..96d4a784 100644 --- a/usecases/web_api_testing/simple_web_api_testing.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py @@ -1,20 +1,21 @@ import time -from dataclasses import dataclass, field -from typing import List, Any, Union, Dict +from dataclasses import dataclass, field from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage -from capabilities import Capability -from capabilities.capability import capabilities_to_action_model -from capabilities.http_request import HTTPRequest -from capabilities.record_note import RecordNote -from capabilities.submit_flag import SubmitFlag -from usecases.common_patterns import RoundBasedUseCase -from usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy -from utils import LLMResult, tool_message, ui -from utils.configurable import parameter -from utils.openai.openai_lib import OpenAILib from rich.panel import Panel -from usecases import use_case +from typing import List, Any, Union, Dict + +from hackingBuddyGPT.capabilities import Capability +from hackingBuddyGPT.capabilities.capability import capabilities_to_action_model +from hackingBuddyGPT.capabilities.http_request import HTTPRequest +from hackingBuddyGPT.capabilities.record_note import RecordNote +from hackingBuddyGPT.capabilities.submit_flag import SubmitFlag +from hackingBuddyGPT.usecases.common_patterns import RoundBasedUseCase +from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy +from hackingBuddyGPT.utils import LLMResult, tool_message, ui +from hackingBuddyGPT.utils.configurable import parameter +from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib +from hackingBuddyGPT.usecases import use_case import pydantic_core diff --git a/utils/__init__.py b/src/hackingBuddyGPT/utils/__init__.py similarity index 86% rename from utils/__init__.py rename to src/hackingBuddyGPT/utils/__init__.py index 6d2870ed..7df80e5e 100644 --- a/utils/__init__.py +++ b/src/hackingBuddyGPT/utils/__init__.py @@ -6,4 +6,4 @@ from .db_storage import * from .openai import * from .psexec import * -from .ssh_connection import * +from .ssh_connection import * \ No newline at end of file diff --git a/utils/cli_history.py b/src/hackingBuddyGPT/utils/cli_history.py similarity index 92% rename from utils/cli_history.py rename to src/hackingBuddyGPT/utils/cli_history.py index e97b9e56..ff80443f 100644 --- a/utils/cli_history.py +++ b/src/hackingBuddyGPT/utils/cli_history.py @@ -1,5 +1,4 @@ -from utils.llm_util import LLM, trim_result_front - +from .llm_util import LLM, trim_result_front class SlidingCliHistory: diff --git a/utils/configurable.py b/src/hackingBuddyGPT/utils/configurable.py similarity index 100% rename from utils/configurable.py rename to src/hackingBuddyGPT/utils/configurable.py diff --git a/utils/console/__init__.py b/src/hackingBuddyGPT/utils/console/__init__.py similarity index 100% rename from utils/console/__init__.py rename to src/hackingBuddyGPT/utils/console/__init__.py diff --git a/utils/console/console.py b/src/hackingBuddyGPT/utils/console/console.py similarity index 81% rename from utils/console/console.py rename to src/hackingBuddyGPT/utils/console/console.py index 4d49b925..e48091e1 100644 --- a/utils/console/console.py +++ b/src/hackingBuddyGPT/utils/console/console.py @@ -1,6 +1,6 @@ from rich import console -from utils.configurable import configurable +from hackingBuddyGPT.utils.configurable import configurable @configurable("console", "Console") diff --git a/utils/db_storage/__init__.py b/src/hackingBuddyGPT/utils/db_storage/__init__.py similarity index 100% rename from utils/db_storage/__init__.py rename to src/hackingBuddyGPT/utils/db_storage/__init__.py diff --git a/utils/db_storage/db_storage.py b/src/hackingBuddyGPT/utils/db_storage/db_storage.py similarity index 99% rename from utils/db_storage/db_storage.py rename to src/hackingBuddyGPT/utils/db_storage/db_storage.py index f7c93232..6187cf87 100644 --- a/utils/db_storage/db_storage.py +++ b/src/hackingBuddyGPT/utils/db_storage/db_storage.py @@ -1,6 +1,6 @@ import sqlite3 -from utils.configurable import configurable, parameter +from hackingBuddyGPT.utils.configurable import configurable, parameter @configurable("db_storage", "Stores the results of the experiments in a SQLite database") diff --git a/utils/llm_util.py b/src/hackingBuddyGPT/utils/llm_util.py similarity index 99% rename from utils/llm_util.py rename to src/hackingBuddyGPT/utils/llm_util.py index 20934a84..658abe44 100644 --- a/utils/llm_util.py +++ b/src/hackingBuddyGPT/utils/llm_util.py @@ -8,7 +8,6 @@ SAFETY_MARGIN = 128 STEP_CUT_TOKENS = 128 - @dataclass class LLMResult: result: typing.Any diff --git a/utils/openai/__init__.py b/src/hackingBuddyGPT/utils/openai/__init__.py similarity index 100% rename from utils/openai/__init__.py rename to src/hackingBuddyGPT/utils/openai/__init__.py diff --git a/utils/openai/openai_lib.py b/src/hackingBuddyGPT/utils/openai/openai_lib.py similarity index 94% rename from utils/openai/openai_lib.py rename to src/hackingBuddyGPT/utils/openai/openai_lib.py index 18df1ad8..f91c484e 100644 --- a/utils/openai/openai_lib.py +++ b/src/hackingBuddyGPT/utils/openai/openai_lib.py @@ -1,13 +1,11 @@ -import time -from dataclasses import dataclass - import instructor -import tiktoken - -from utils import LLM, configurable, LLMResult import openai +import tiktoken +import time +from dataclasses import dataclass -from utils.configurable import parameter +from hackingBuddyGPT.utils import LLM, configurable, LLMResult +from hackingBuddyGPT.utils.configurable import parameter @configurable("openai-lib", "OpenAI Library based connection") diff --git a/utils/openai/openai_llm.py b/src/hackingBuddyGPT/utils/openai/openai_llm.py similarity index 96% rename from utils/openai/openai_llm.py rename to src/hackingBuddyGPT/utils/openai/openai_llm.py index eb792111..8184057f 100644 --- a/utils/openai/openai_llm.py +++ b/src/hackingBuddyGPT/utils/openai/openai_llm.py @@ -1,13 +1,11 @@ -import time - import requests +import tiktoken +import time from dataclasses import dataclass -import tiktoken - -from utils.configurable import configurable, parameter -from utils.llm_util import LLMResult, LLM +from hackingBuddyGPT.utils.configurable import configurable, parameter +from hackingBuddyGPT.utils.llm_util import LLMResult, LLM @configurable("openai-compatible-llm-api", "OpenAI-compatible LLM API") @dataclass diff --git a/utils/psexec/__init__.py b/src/hackingBuddyGPT/utils/psexec/__init__.py similarity index 100% rename from utils/psexec/__init__.py rename to src/hackingBuddyGPT/utils/psexec/__init__.py diff --git a/utils/psexec/psexec.py b/src/hackingBuddyGPT/utils/psexec/psexec.py similarity index 94% rename from utils/psexec/psexec.py rename to src/hackingBuddyGPT/utils/psexec/psexec.py index 54f19aea..dcc95240 100644 --- a/utils/psexec/psexec.py +++ b/src/hackingBuddyGPT/utils/psexec/psexec.py @@ -1,10 +1,8 @@ -import warnings from dataclasses import dataclass -from typing import Tuple - from pypsexec.client import Client +from typing import Tuple -from utils.configurable import configurable +from hackingBuddyGPT.utils.configurable import configurable @configurable("psexec", "connects to a remote host via PSExec") diff --git a/src/hackingBuddyGPT/utils/shell_root_detection.py b/src/hackingBuddyGPT/utils/shell_root_detection.py new file mode 100644 index 00000000..318e8434 --- /dev/null +++ b/src/hackingBuddyGPT/utils/shell_root_detection.py @@ -0,0 +1,14 @@ +import re + +GOT_ROOT_REGEXPs = [ + re.compile("^# $"), + re.compile("^bash-[0-9]+.[0-9]# $") +] + +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 \ No newline at end of file diff --git a/utils/ssh_connection/__init__.py b/src/hackingBuddyGPT/utils/ssh_connection/__init__.py similarity index 100% rename from utils/ssh_connection/__init__.py rename to src/hackingBuddyGPT/utils/ssh_connection/__init__.py diff --git a/utils/ssh_connection/ssh_connection.py b/src/hackingBuddyGPT/utils/ssh_connection/ssh_connection.py similarity index 95% rename from utils/ssh_connection/ssh_connection.py rename to src/hackingBuddyGPT/utils/ssh_connection/ssh_connection.py index dcf0a848..33bf8557 100644 --- a/utils/ssh_connection/ssh_connection.py +++ b/src/hackingBuddyGPT/utils/ssh_connection/ssh_connection.py @@ -3,7 +3,7 @@ from fabric import Connection from typing import Optional, Tuple -from utils .configurable import configurable +from hackingBuddyGPT.utils.configurable import configurable @configurable("ssh", "connects to a remote host via SSH") diff --git a/utils/ui.py b/src/hackingBuddyGPT/utils/ui.py similarity index 94% rename from utils/ui.py rename to src/hackingBuddyGPT/utils/ui.py index 7934245f..753ec223 100644 --- a/utils/ui.py +++ b/src/hackingBuddyGPT/utils/ui.py @@ -1,7 +1,6 @@ from rich.table import Table -from utils.db_storage.db_storage import DbStorage - +from .db_storage.db_storage import DbStorage # helper to fill the history table with data from the db def get_history_table(enable_explanation: bool, enable_update_state: bool, run_id: int, db: DbStorage, turn: int) -> Table: diff --git a/tests/.gitkeep b/tests/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_root_detection.py b/tests/test_root_detection.py new file mode 100644 index 00000000..9c3e94e2 --- /dev/null +++ b/tests/test_root_detection.py @@ -0,0 +1,7 @@ +from hackingBuddyGPT.utils.shell_root_detection import got_root + +def test_got_root(): + hostname = "i_dont_care" + + assert got_root(hostname, "# ") == True + assert got_root(hostname, "$ ") == False