move openapi parsing into utils

Author: andreashappe

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -1,16 +1,14 @@
 import os
 from dataclasses import field
-from typing import Dict
 
 from rich.panel import Panel
 
-from hackingBuddyGPT.capabilities import Capability
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
 from hackingBuddyGPT.capabilities.record_note import RecordNote
-from hackingBuddyGPT.usecases.agents import Agent
-from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case
+from hackingBuddyGPT.usecases.base import AutonomousUseCase, use_case
 from hackingBuddyGPT.usecases.web_api_testing.documentation.openapi_specification_handler import \
     OpenAPISpecificationHandler
+from hackingBuddyGPT.utils.capability_manager import CapabilityManager
 from hackingBuddyGPT.utils.prompt_generation import PromptGenerationHelper
 from hackingBuddyGPT.utils.prompt_generation.information import PromptContext
 from hackingBuddyGPT.utils.prompt_generation.prompt_engineer import PromptEngineer
@@ -23,7 +21,8 @@
 from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 
 
-class SimpleWebAPIDocumentation(Agent):
+@use_case("Minimal implementation of a web API testing use case")
+class SimpleWebAPIDocumentation(AutonomousUseCase):
     """
          SimpleWebAPIDocumentation is an agent class for automating REST API documentation.
 
@@ -41,10 +40,10 @@ class SimpleWebAPIDocumentation(Agent):
             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
+    llm: OpenAILib = None
     _prompt_history: Prompt = field(default_factory=list)
     _context: Context = field(default_factory=lambda: {"notes": list()})
-    _capabilities: Dict[str, Capability] = field(default_factory=dict)
+    _capabilities: CapabilityManager = None
     _all_http_methods_found: bool = False
     config_path: str = parameter(
         desc="Configuration file path",
@@ -75,6 +74,8 @@ class SimpleWebAPIDocumentation(Agent):
         default="GET,POST,PUT,PATCH,DELETE",
     )
 
+    def get_name(self) -> str:
+        return self.__class__.__name__
 
     def init(self):
         """Initialize the agent with configurations, capabilities, and handlers."""
@@ -90,7 +91,11 @@ def init(self):
 
         self.categorized_endpoints = self.categorize_endpoints(self._correct_endpoints, query_params)
 
-        self._setup_capabilities()
+        # setup capabilities
+        self._capabilities.init()
+        self._capabilities.add_capability(HTTPRequest(self.host))
+        self._capabilities.add_capability(RecordNote(self._context["notes"]))
+
         self._prompt_context = PromptContext.DOCUMENTATION
         name, initial_prompt = self._setup_initial_prompt(description=description)
         self._initialize_handlers(config=config, description=description, token=token, name=name,
@@ -150,7 +155,7 @@ def _initialize_handlers(self, config, description, token, name, initial_prompt)
            """
         self.all_capabilities = {
                                  "http_request": HTTPRequest(self.host)}
-        self._llm_handler = LLMHandler(self.llm, self._capabilities,  all_possible_capabilities=self.all_capabilities)
+        self._llm_handler = LLMHandler(self.llm, self._capabilities._capabilities,  all_possible_capabilities=self.all_capabilities)
 
         self._response_handler = ResponseHandler(llm_handler=self._llm_handler, prompt_context=self._prompt_context,
                                                  prompt_helper=self.prompt_helper, config=config)
@@ -226,25 +231,6 @@ def categorize_endpoints(self, endpoints, query: dict):
             "multi-level_resource": multi_level_resource,
         }
 
-
-
-    def _setup_capabilities(self):
-        """
-           Initializes the LLM agent's capabilities for interacting with the API.
-
-           This sets up tool wrappers that the language model can call, such as:
-           - `http_request`: For performing HTTP calls against the target API.
-           - `record_note`: For storing observations, notes, or documentation artifacts.
-
-           Side Effects:
-               - Populates `self._capabilities` with callable tools used during exploration and documentation.
-           """
-        """Initializes agent's capabilities for API documentation."""
-        self._capabilities = {
-            "http_request": HTTPRequest(self.host),
-            "record_note": RecordNote(self._context["notes"])
-        }
-
     def all_http_methods_found(self, turn: int) -> bool:
         """
            Checks whether all expected HTTP methods (GET, POST, PUT, DELETE) have been discovered
@@ -441,10 +427,4 @@ def run_documentation(self, turn: int, move_type: str) -> None:
             self._evaluator.finalize_documentation_metrics(
                 file_path=self._documentation_handler.file.split(".yaml")[0] + ".txt")
 
-        self.all_http_methods_found(turn)
-
-
-@use_case("Minimal implementation of a web API testing use case")
-class SimpleWebAPIDocumentationUseCase(AutonomousAgentUseCase[SimpleWebAPIDocumentation]):
-    """Use case for the SimpleWebAPIDocumentation agent."""
-    pass
+        self.all_http_methods_found(turn)

src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py

@@ -8,17 +8,16 @@
 import pydantic_core
 from rich.panel import Panel
 
-from hackingBuddyGPT.capabilities import Capability
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
 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.agents import Agent
-from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case
+from hackingBuddyGPT.usecases.base import AutonomousUseCase, use_case
+from hackingBuddyGPT.utils.capability_manager import CapabilityManager
 from hackingBuddyGPT.utils.prompt_generation import PromptGenerationHelper
 from hackingBuddyGPT.utils.prompt_generation.information import PenTestingInformation
 from hackingBuddyGPT.utils.prompt_generation.information import PromptPurpose
-from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing import OpenAPISpecificationParser
+from hackingBuddyGPT.utils.openapi.openapi_parser import OpenAPISpecificationParser
 from hackingBuddyGPT.usecases.web_api_testing.documentation.report_handler import ReportHandler
 from hackingBuddyGPT.utils.prompt_generation.information import PromptContext
 from hackingBuddyGPT.utils.prompt_generation.prompt_engineer import PromptEngineer
@@ -36,8 +35,8 @@
 
 # OpenAPI specification file path
 
-
-class SimpleWebAPITesting(Agent):
+@use_case("Minimal implementation of a web API testing use case")
+class SimpleWebAPITesting(AutonomousUseCase):
     """
     SimpleWebAPITesting is an agent class for automating web API testing.
 
@@ -53,7 +52,7 @@ class SimpleWebAPITesting(Agent):
         _all_test_cases_run (bool): Flag indicating if all HTTP methods have been found.
     """
 
-    llm: OpenAILib
+    llm: OpenAILib = None
     host: str = parameter(desc="The host to test", default="https://jsonplaceholder.typicode.com")
     config_path: str = parameter(
         desc="Configuration file path",
@@ -71,7 +70,7 @@ class SimpleWebAPITesting(Agent):
     )
     _prompt_history: Prompt = field(default_factory=list)
     _context: Context = field(default_factory=lambda: {"notes": list(), "test_cases": list(), "parsed": list()})
-    _capabilities: Dict[str, Capability] = field(default_factory=dict)
+    _capabilities: CapabilityManager = None
     _all_test_cases_run: bool = False
 
     def init(self):
@@ -86,6 +85,9 @@ def init(self):
         self._setup_initial_prompt()
         self.last_prompt = ""
 
+    def get_name(self) -> str:
+        return self.__class__.__name__
+
     def _load_openapi_specification(self):
         """
            Loads the OpenAPI specification from the configured file path.
@@ -108,6 +110,11 @@ def _setup_environment(self):
            - Setting the prompt context to `PromptContext.PENTESTING`.
            """
         self._context["host"] = self.host
+
+        # setup capabilities
+        self._capabilities = CapabilityManager(self.log)
+        self._capabilities.add_capability(HTTPRequest(self.host))
+
         self._setup_capabilities()
         self.categorized_endpoints = self._openapi_specification_parser.categorize_endpoints(self.correct_endpoints,
                                                                                              self.query_params)
@@ -128,7 +135,7 @@ def _setup_handlers(self):
 
             If username and password are not found in the config, defaults are used.
             """
-        self._llm_handler = LLMHandler(self.llm, self._capabilities, all_possible_capabilities=self.all_capabilities)
+        self._llm_handler = LLMHandler(self.llm, self._capabilities._capabilities, all_possible_capabilities=self.all_capabilities)
         self.prompt_helper = PromptGenerationHelper(self.host, self.description)
         if "username" in self.config.keys() and "password" in self.config.keys():
             username = self.config.get("username")
@@ -194,8 +201,6 @@ def _setup_capabilities(self) -> None:
         test_cases = self._context["test_cases"]
         self.python_test_case_capability = {"python_test_case": PythonTestCase(test_cases)}
         self.parse_capacity = {"parse": ParsedInformation(test_cases)}
-        self._capabilities = {
-            "http_request": HTTPRequest(self.host)}
         self.all_capabilities = {"python_test_case": PythonTestCase(test_cases), "parse": ParsedInformation(test_cases),
                                  "http_request": HTTPRequest(self.host),
                                  "record_note": RecordNote(notes)}
@@ -539,14 +544,4 @@ def execute_response(self, response, completion):
         self._prompt_history.append(tool_message(self._response_handler.extract_key_elements_of_response(result), tool_call_id))
 
         self.adjust_user(result)
-        return result
-
-
-@use_case("Minimal implementation of a web API testing use case")
-class SimpleWebAPITestingUseCase(AutonomousAgentUseCase[SimpleWebAPITesting]):
-    """
-    A use case for the SimpleWebAPITesting agent, encapsulating the setup and execution
-    of the web API testing scenario.
-    """
-
-    pass
+        return result

src/hackingBuddyGPT/usecases/web_api_testing/utils/endpoint_categorizer.py

@@ -36,7 +36,6 @@ def run(self, recorded_note: str) -> None:
             The current implementation is commented out and serves as a placeholder for
             integrating with OpenAI's API. Uncomment and modify the code as needed.
         """
-        """
         assistant = self.client.beta.assistants.create(
             name="Yaml File Analysis Assistant",
             instructions="You are an OpenAPI specification analyst. Use your knowledge to check "
@@ -86,5 +85,4 @@ def run(self, recorded_note: str) -> None:
         )
 
         # The thread now has a vector store with that file in its tool resources.
-        print(thread.tool_resources.file_search)
-        """
+        print(thread.tool_resources.file_search)

src/hackingBuddyGPT/usecases/web_api_testing/documentation/parsing/__init__.py renamed to src/hackingBuddyGPT/utils/openapi/__init__.py

@@ -9,7 +9,7 @@
 
 import pandas
 
-from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing import OpenAPISpecificationParser
+from hackingBuddyGPT.utils.openapi.openapi_parser import OpenAPISpecificationParser
 from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import (
     PromptPurpose,
 )