diff --git a/pyproject.toml b/pyproject.toml index 634565f..52c2586 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,6 @@ dependencies = [ 'nltk', 'fastapi == 0.114.0', 'fastapi-utils == 0.7.0', - 'jinja2 == 3.1.6', 'uvicorn[standard] == 0.30.6', 'dataclasses_json == 0.6.7', 'websockets == 13.1', diff --git a/src/hackingBuddyGPT/cli/wintermute.py b/src/hackingBuddyGPT/cli/wintermute.py index fef6095..94252eb 100644 --- a/src/hackingBuddyGPT/cli/wintermute.py +++ b/src/hackingBuddyGPT/cli/wintermute.py @@ -1,7 +1,7 @@ import argparse import sys -from hackingBuddyGPT.usecases.base import use_cases +from hackingBuddyGPT.usecases.usecase import use_cases from hackingBuddyGPT.utils.configurable import CommandMap, InvalidCommand, Parseable, instantiate diff --git a/src/hackingBuddyGPT/strategies.py b/src/hackingBuddyGPT/strategies.py index c1587a1..655b952 100644 --- a/src/hackingBuddyGPT/strategies.py +++ b/src/hackingBuddyGPT/strategies.py @@ -4,9 +4,10 @@ from dataclasses import dataclass from mako.template import Template from hackingBuddyGPT.capability import capabilities_to_simple_text_handler -from hackingBuddyGPT.usecases.base import UseCase +from hackingBuddyGPT.usecases.usecase import UseCase from hackingBuddyGPT.utils import llm_util from hackingBuddyGPT.utils.histories import HistoryCmdOnly, HistoryFull, HistoryNone +from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection from hackingBuddyGPT.utils.logging import log_conversation, Logger, log_param, log_section from hackingBuddyGPT.utils.capability_manager import CapabilityManager @@ -140,3 +141,59 @@ def check_success(self, cmd:str, result:str) -> bool: def postprocess_commands(self, cmd:str) -> List[str]: return [cmd] + +@dataclass +class SimpleStrategy(UseCase, abc.ABC): + max_turns: int = 10 + + llm: OpenAILib = None + + log: Logger = log_param + + _got_root: bool = False + + _capabilities: CapabilityManager = None + + def init(self): + super().init() + self._capabilities = CapabilityManager(self.log) + + @abc.abstractmethod + def perform_round(self, turn: int): + pass + + def before_run(self): + pass + + def after_run(self): + pass + + def run(self, configuration): + self.configuration = configuration + self.log.start_run(self.get_name(), self.serialize_configuration(configuration)) + + self.before_run() + + turn = 1 + try: + while turn <= self.max_turns and not self._got_root: + with self.log.section(f"round {turn}"): + self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}") + + self._got_root = self.perform_round(turn) + + turn += 1 + + self.after_run() + + # write the final result to the database and console + if self._got_root: + self.log.run_was_success() + else: + self.log.run_was_failure("maximum turn number reached") + + return self._got_root + except Exception: + import traceback + self.log.run_was_failure("exception occurred", details=f":\n\n{traceback.format_exc()}") + raise \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/agents.py b/src/hackingBuddyGPT/usecases/agents.py deleted file mode 100644 index 6c2996b..0000000 --- a/src/hackingBuddyGPT/usecases/agents.py +++ /dev/null @@ -1,124 +0,0 @@ -import datetime -from abc import ABC, abstractmethod -from dataclasses import dataclass, field -from mako.template import Template -from typing import Dict - -from hackingBuddyGPT.utils.logging import log_conversation, Logger, log_param -from hackingBuddyGPT.capability import ( - Capability, - capabilities_to_simple_text_handler, -) -from hackingBuddyGPT.utils import llm_util -from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection - - -@dataclass -class Agent(ABC): - log: Logger = log_param - - _capabilities: Dict[str, Capability] = field(default_factory=dict) - _default_capability: Capability = None - - llm: OpenAIConnection = None - - def init(self): # noqa: B027 - pass - - def before_run(self): # noqa: B027 - pass - - def after_run(self): # noqa: B027 - pass - - # callback - @abstractmethod - def perform_round(self, turn: int) -> bool: - pass - - def add_capability(self, cap: Capability, name: str = None, default: bool = False): - if name is None: - name = cap.get_name() - self._capabilities[name] = cap - if default: - self._default_capability = cap - - def get_capability(self, name: str) -> Capability: - return self._capabilities.get(name, self._default_capability) - - def run_capability_json(self, message_id: int, tool_call_id: str, capability_name: str, arguments: str) -> str: - capability = self.get_capability(capability_name) - - tic = datetime.datetime.now() - try: - result = capability.to_model().model_validate_json(arguments).execute() - except Exception as e: - result = f"EXCEPTION: {e}" - duration = datetime.datetime.now() - tic - - self.log.add_tool_call(message_id, tool_call_id, capability_name, arguments, result, duration) - return result - - def run_capability_simple_text(self, message_id: int, cmd: str) -> tuple[str, str, str, bool]: - _capability_descriptions, parser = capabilities_to_simple_text_handler(self._capabilities, default_capability=self._default_capability) - - tic = datetime.datetime.now() - try: - success, output = parser(cmd) - except Exception as e: - success = False - output = f"EXCEPTION: {e}" - duration = datetime.datetime.now() - tic - - if not success: - self.log.add_tool_call(message_id, tool_call_id=0, function_name="", arguments=cmd, result_text=output[0], duration=0) - return "", "", output, False - - capability, cmd, (result, got_root) = output - self.log.add_tool_call(message_id, tool_call_id=0, function_name=capability, arguments=cmd, result_text=result, duration=duration) - - return capability, cmd, result, got_root - - def get_capability_block(self) -> str: - capability_descriptions, _parser = capabilities_to_simple_text_handler(self._capabilities) - return "You can either\n\n" + "\n".join(f"- {description}" for description in capability_descriptions.values()) - - -@dataclass -class AgentWorldview(ABC): - @abstractmethod - def to_template(self): - pass - - @abstractmethod - def update(self, capability, cmd, result): - pass - - -class TemplatedAgent(Agent): - _state: AgentWorldview = None - _template: Template = None - _template_size: int = 0 - - def init(self): - super().init() - - def set_initial_state(self, initial_state: AgentWorldview): - self._state = initial_state - - def set_template(self, template: str): - self._template = Template(filename=template) - self._template_size = self.llm.count_tokens(self._template.source) - - @log_conversation("Asking LLM for a new command...") - def perform_round(self, turn: int) -> bool: - # get the next command from the LLM - answer = self.llm.get_response(self._template, capabilities=self.get_capability_block(), **self._state.to_template()) - message_id = self.log.call_response(answer) - - capability, cmd, result, got_root = self.run_capability_simple_text(message_id, llm_util.cmd_output_fixer(answer.result)) - - self._state.update(capability, cmd, result) - - # if we got root, we can stop the loop - return got_root diff --git a/src/hackingBuddyGPT/usecases/base.py b/src/hackingBuddyGPT/usecases/base.py deleted file mode 100644 index 9f1896e..0000000 --- a/src/hackingBuddyGPT/usecases/base.py +++ /dev/null @@ -1,162 +0,0 @@ -import abc -import json -import argparse -from dataclasses import dataclass - -from hackingBuddyGPT.utils.logging import Logger, log_param -from typing import Dict, Type, TypeVar, Generic - -from hackingBuddyGPT.utils.configurable import Transparent, configurable - -@dataclass -class UseCase(abc.ABC): - """ - A UseCase is the combination of tools and capabilities to solve a specific problem. - It is usually recommended, to have a UseCase be a dataclass, with all the necessary utils (being of type - @configurable) as fields. Then they can be automatically injected from the command line / environment / .env - parameters. - - All UseCases should inherit from this class, implement the run method, and be decorated with the @use_case decorator, - so that they can be automatically discovered and run from the command line. - """ - - log: Logger = log_param - - 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. - """ - pass - - def serialize_configuration(self, configuration) -> str: - return json.dumps(configuration) - - @abc.abstractmethod - 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 - functionalities of each step. - """ - pass - - @abc.abstractmethod - def get_name(self) -> str: - """ - This method should return the name of the use case. It is used for logging and debugging purposes. - """ - pass - - -# this runs the main loop for a bounded amount of turns or until root was achieved -@dataclass -class AutonomousUseCase(UseCase, abc.ABC): - max_turns: int = 10 - - _got_root: bool = False - - @abc.abstractmethod - def perform_round(self, turn: int): - pass - - def before_run(self): - pass - - def after_run(self): - pass - - def run(self, configuration): - self.configuration = configuration - self.log.start_run(self.get_name(), self.serialize_configuration(configuration)) - - self.before_run() - - turn = 1 - try: - while turn <= self.max_turns and not self._got_root: - with self.log.section(f"round {turn}"): - self.log.console.log(f"[yellow]Starting turn {turn} of {self.max_turns}") - - self._got_root = self.perform_round(turn) - - turn += 1 - - self.after_run() - - # write the final result to the database and console - if self._got_root: - self.log.run_was_success() - else: - self.log.run_was_failure("maximum turn number reached") - - return self._got_root - except Exception: - import traceback - self.log.run_was_failure("exception occurred", details=f":\n\n{traceback.format_exc()}") - raise - - -use_cases: Dict[str, configurable] = dict() - - -T = TypeVar("T", bound=type) - - -class AutonomousAgentUseCase(AutonomousUseCase, Generic[T]): - agent: T = None - - def perform_round(self, turn: int): - raise ValueError("Do not use AutonomousAgentUseCase without supplying an agent type as generic") - - def get_name(self) -> str: - raise ValueError("Do not use AutonomousAgentUseCase without supplying an agent type as generic") - - @classmethod - def __class_getitem__(cls, item): - item = dataclass(item) - - class AutonomousAgentUseCase(AutonomousUseCase): - agent: Transparent(item) = None - - def init(self): - super().init() - self.agent.init() - - def get_name(self) -> str: - return self.__class__.__name__ - - def before_run(self): - return self.agent.before_run() - - def after_run(self): - return self.agent.after_run() - - def perform_round(self, turn: int): - return self.agent.perform_round(turn) - - constructed_class = dataclass(AutonomousAgentUseCase) - - return constructed_class - - -def use_case(description): - def inner(cls): - cls = dataclass(cls) - name = cls.__name__.removesuffix("UseCase") - if name in use_cases: - raise IndexError(f"Use case with name {name} already exists") - use_cases[name] = configurable(name, description)(cls) - return cls - - return inner - - -def register_use_case(name: str, description: str, use_case: Type[UseCase]): - """ - This function is used to register a UseCase that was created manually, and not through the use_case decorator. - """ - if name in use_cases: - raise IndexError(f"Use case with name {name} already exists") - use_cases[name] = configurable(name, description)(use_case) diff --git a/src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py b/src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py index b7f8a6d..af84469 100644 --- a/src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py +++ b/src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py @@ -1,7 +1,7 @@ from mako.template import Template from hackingBuddyGPT.capabilities import SSHRunCommand -from hackingBuddyGPT.usecases.base import UseCase, use_case +from hackingBuddyGPT.usecases.usecase import UseCase, use_case from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection diff --git a/src/hackingBuddyGPT/usecases/linux_privesc.py b/src/hackingBuddyGPT/usecases/linux_privesc.py index 994a564..f2f2d71 100644 --- a/src/hackingBuddyGPT/usecases/linux_privesc.py +++ b/src/hackingBuddyGPT/usecases/linux_privesc.py @@ -6,7 +6,7 @@ from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential from hackingBuddyGPT.strategies import CommandStrategy -from hackingBuddyGPT.usecases.base import use_case +from hackingBuddyGPT.usecases.usecase import use_case from hackingBuddyGPT.utils import llm_util from hackingBuddyGPT.utils.logging import log_conversation from hackingBuddyGPT.utils.rag import RagBackground diff --git a/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py b/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py index 70ff27b..d26268a 100644 --- a/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py +++ b/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py @@ -1,7 +1,7 @@ import re from typing import List from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential -from hackingBuddyGPT.usecases.base import use_case +from hackingBuddyGPT.usecases.usecase import use_case from hackingBuddyGPT.strategies import CommandStrategy from hackingBuddyGPT.utils import llm_util from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection diff --git a/src/hackingBuddyGPT/usecases/usecase.py b/src/hackingBuddyGPT/usecases/usecase.py new file mode 100644 index 0000000..505dd13 --- /dev/null +++ b/src/hackingBuddyGPT/usecases/usecase.py @@ -0,0 +1,71 @@ +import abc +import json +from dataclasses import dataclass + +from hackingBuddyGPT.utils.logging import Logger, log_param +from typing import Dict, Type + +from hackingBuddyGPT.utils.configurable import configurable + +@dataclass +class UseCase(abc.ABC): + """ + A UseCase is the combination of tools and capabilities to solve a specific problem. + It is usually recommended, to have a UseCase be a dataclass, with all the necessary utils (being of type + @configurable) as fields. Then they can be automatically injected from the command line / environment / .env + parameters. + + All UseCases should inherit from this class, implement the run method, and be decorated with the @use_case decorator, + so that they can be automatically discovered and run from the command line. + """ + + log: Logger = log_param + + 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. + """ + pass + + def serialize_configuration(self, configuration) -> str: + return json.dumps(configuration) + + @abc.abstractmethod + 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 + functionalities of each step. + """ + pass + + @abc.abstractmethod + def get_name(self) -> str: + """ + This method should return the name of the use case. It is used for logging and debugging purposes. + """ + pass + +use_cases: Dict[str, configurable] = dict() + +def use_case(description): + def inner(cls): + cls = dataclass(cls) + name = cls.__name__.removesuffix("UseCase") + if name in use_cases: + raise IndexError(f"Use case with name {name} already exists") + use_cases[name] = configurable(name, description)(cls) + return cls + + return inner + + +def register_use_case(name: str, description: str, use_case: Type[UseCase]): + """ + This function is used to register a UseCase that was created manually, and not through the use_case decorator. + """ + if name in use_cases: + raise IndexError(f"Use case with name {name} already exists") + use_cases[name] = configurable(name, description)(use_case) diff --git a/src/hackingBuddyGPT/usecases/viewer.py b/src/hackingBuddyGPT/usecases/viewer.py index b4da563..b4e9f8e 100644 --- a/src/hackingBuddyGPT/usecases/viewer.py +++ b/src/hackingBuddyGPT/usecases/viewer.py @@ -17,7 +17,7 @@ from starlette.staticfiles import StaticFiles from starlette.templating import Jinja2Templates -from hackingBuddyGPT.usecases.base import UseCase, use_case +from hackingBuddyGPT.usecases.usecase 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 ( diff --git a/src/hackingBuddyGPT/usecases/web/with_explanation.py b/src/hackingBuddyGPT/usecases/web/with_explanation.py index 5ed7f44..4b99400 100644 --- a/src/hackingBuddyGPT/usecases/web/with_explanation.py +++ b/src/hackingBuddyGPT/usecases/web/with_explanation.py @@ -1,24 +1,20 @@ from dataclasses import field -from typing import List, Any, Union, Dict, Iterable, Optional - +from typing import List, Any, Union, Iterable, Optional from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage from openai.types.chat.chat_completion_chunk import ChoiceDelta -from hackingBuddyGPT.capabilities import Capability from hackingBuddyGPT.capabilities.http_request import HTTPRequest from hackingBuddyGPT.capabilities.submit_flag import SubmitFlag -from hackingBuddyGPT.usecases.agents import Agent -from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case +from hackingBuddyGPT.strategies import SimpleStrategy +from hackingBuddyGPT.usecases.usecase import use_case from hackingBuddyGPT.utils import LLMResult, tool_message from hackingBuddyGPT.utils.configurable import parameter -from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib Prompt = List[Union[ChatCompletionMessage, ChatCompletionMessageParam]] Context = Any - -class WebTestingWithExplanation(Agent): - llm: OpenAILib +@use_case("Minimal implementation of a web testing use case while allowing the llm to 'talk'") +class WebTestingWithExplanation(SimpleStrategy): host: str = parameter(desc="The host to test", default="http://localhost") flag_format_description: str = parameter( desc="Description of the flag provided to the LLM", @@ -35,14 +31,16 @@ class WebTestingWithExplanation(Agent): _prompt_history: Prompt = field(default_factory=list) _context: Context = field(default_factory=lambda: {"notes": list()}) - _capabilities: Dict[str, Capability] = field(default_factory=dict) _all_flags_found: bool = False def init(self): super().init() self._context["host"] = self.host - self.add_capability(SubmitFlag(self.flag_format_description, set(self.flag_template.format(flag=flag) for flag in self.flags.split(",")), success_function=self.all_flags_found)) - self.add_capability(HTTPRequest(self.host)) + self._capabilities.add_capability(SubmitFlag(self.flag_format_description, set(self.flag_template.format(flag=flag) for flag in self.flags.split(",")), success_function=self.all_flags_found)) + self._capabilities.add_capability(HTTPRequest(self.host)) + + def get_name(self) -> str: + return self.__class__.__name__ def before_run(self): system_message = ( @@ -64,7 +62,7 @@ def all_flags_found(self): def perform_round(self, turn: int): prompt = self._prompt_history # TODO: in the future, this should do some context truncation - result_stream: Iterable[Union[ChoiceDelta, LLMResult]] = self.llm.stream_response(prompt, self.log.console, capabilities=self._capabilities, get_individual_updates=True) + result_stream: Iterable[Union[ChoiceDelta, LLMResult]] = self.llm.stream_response(prompt, self.log.console, capabilities=self._capabilities._capabilities, get_individual_updates=True) result: Optional[LLMResult] = None stream_output = self.log.stream_message("assistant") # TODO: do not hardcode the role for delta in result_stream: @@ -83,12 +81,7 @@ def perform_round(self, turn: int): if message.tool_calls is not None: for tool_call in message.tool_calls: - tool_result = self.run_capability_json(message_id, tool_call.id, tool_call.function.name, tool_call.function.arguments) + tool_result = self._capabilities.run_capability_json(message_id, tool_call.id, tool_call.function.name, tool_call.function.arguments) self._prompt_history.append(tool_message(tool_result, tool_call.id)) return self._all_flags_found - - -@use_case("Minimal implementation of a web testing use case while allowing the llm to 'talk'") -class WebTestingWithExplanationUseCase(AutonomousAgentUseCase[WebTestingWithExplanation]): - pass diff --git a/src/hackingBuddyGPT/usecases/web_api_documentation/simple_openapi_documentation.py b/src/hackingBuddyGPT/usecases/web_api_documentation/simple_openapi_documentation.py index 77f6a36..085c945 100644 --- a/src/hackingBuddyGPT/usecases/web_api_documentation/simple_openapi_documentation.py +++ b/src/hackingBuddyGPT/usecases/web_api_documentation/simple_openapi_documentation.py @@ -1,17 +1,13 @@ import json -from logging import config import os -from dataclasses import field - -from rich.panel import Panel +from dataclasses import field from hackingBuddyGPT.capabilities.http_request import HTTPRequest from hackingBuddyGPT.capabilities.record_note import RecordNote -from hackingBuddyGPT.usecases.base import AutonomousUseCase, use_case +from hackingBuddyGPT.strategies import SimpleStrategy +from hackingBuddyGPT.usecases.usecase import use_case from hackingBuddyGPT.usecases.web_api_documentation.openapi_specification_handler import \ OpenAPISpecificationHandler -from hackingBuddyGPT.utils.capability_manager import CapabilityManager -from hackingBuddyGPT.utils.logging import Logger, log_param from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import PromptStrategy from hackingBuddyGPT.utils.prompt_generation.prompt_generation_helper import PromptGenerationHelper from hackingBuddyGPT.utils.prompt_generation.information import PromptContext @@ -21,11 +17,11 @@ from hackingBuddyGPT.utils.web_api.custom_datatypes import Context, Prompt from hackingBuddyGPT.usecases.web_api_documentation.evaluator import Evaluator from hackingBuddyGPT.utils.configurable import parameter -from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib +from rich.panel import Panel @use_case("Minimal implementation of a web API testing use case") -class SimpleWebAPIDocumentation(AutonomousUseCase): +class SimpleWebAPIDocumentation(SimpleStrategy): """ SimpleWebAPIDocumentation is an agent class for automating REST API documentation. @@ -42,11 +38,8 @@ class SimpleWebAPIDocumentation(AutonomousUseCase): found_all_http_methods (bool): Flag indicating whether all HTTP methods have been found. all_steps_done (bool): Flag to indicate whether the full documentation process is complete. """ - llm: OpenAILib = None - log: Logger = log_param _prompt_history: Prompt = field(default_factory=list) _context: Context = field(default_factory=lambda: {"notes": list()}) - _capabilities: CapabilityManager = None _all_http_methods_found: bool = False config_path: str = parameter( desc="Configuration file path", @@ -113,7 +106,6 @@ def init(self): self.categorized_endpoints = self.categorize_endpoints(self._correct_endpoints, query_params) # setup capabilities - self._capabilities = CapabilityManager(self.log) self._capabilities.add_capability(HTTPRequest(self.host)) self._capabilities.add_capability(RecordNote(self._context["notes"])) diff --git a/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py b/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py index 0652e45..fc44ad4 100644 --- a/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py +++ b/src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py @@ -12,8 +12,8 @@ from hackingBuddyGPT.capabilities.parsed_information import ParsedInformation from hackingBuddyGPT.capabilities.python_test_case import PythonTestCase from hackingBuddyGPT.capabilities.record_note import RecordNote -from hackingBuddyGPT.usecases.base import AutonomousUseCase, use_case -from hackingBuddyGPT.utils.capability_manager import CapabilityManager +from hackingBuddyGPT.strategies import SimpleStrategy +from hackingBuddyGPT.usecases.usecase import use_case from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import PromptStrategy from hackingBuddyGPT.utils.prompt_generation.prompt_generation_helper import PromptGenerationHelper from hackingBuddyGPT.utils.prompt_generation.information import PenTestingInformation @@ -30,13 +30,12 @@ from hackingBuddyGPT.utils.web_api.llm_handler import LLMHandler from hackingBuddyGPT.utils import tool_message from hackingBuddyGPT.utils.configurable import parameter -from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib # OpenAPI specification file path @use_case("Minimal implementation of a web API testing use case") -class SimpleWebAPITesting(AutonomousUseCase): +class SimpleWebAPITesting(SimpleStrategy): """ SimpleWebAPITesting is an agent class for automating web API testing. @@ -52,7 +51,6 @@ class SimpleWebAPITesting(AutonomousUseCase): _all_test_cases_run (bool): Flag indicating if all HTTP methods have been found. """ - llm: OpenAILib = None host: str = parameter(desc="The host to test", default="https://jsonplaceholder.typicode.com") config_path: str = parameter( desc="Configuration file path", @@ -70,7 +68,6 @@ class SimpleWebAPITesting(AutonomousUseCase): ) _prompt_history: Prompt = field(default_factory=list) _context: Context = field(default_factory=lambda: {"notes": list(), "test_cases": list(), "parsed": list()}) - _capabilities: CapabilityManager = None _all_test_cases_run: bool = False def get_strategy(self, strategy_string): @@ -132,7 +129,6 @@ def _setup_environment(self): self._context["host"] = self.host # setup capabilities - self._capabilities = CapabilityManager(self.log) self._capabilities.add_capability(HTTPRequest(self.host)) self._setup_capabilities() diff --git a/src/hackingBuddyGPT/usecases/windows_privesc.py b/src/hackingBuddyGPT/usecases/windows_privesc.py index 537e967..5290a3c 100644 --- a/src/hackingBuddyGPT/usecases/windows_privesc.py +++ b/src/hackingBuddyGPT/usecases/windows_privesc.py @@ -1,4 +1,4 @@ -from hackingBuddyGPT.usecases.base import use_case +from hackingBuddyGPT.usecases.usecase import use_case from hackingBuddyGPT.strategies import CommandStrategy from hackingBuddyGPT.utils import SSHConnection