Refactor code

src/hackingBuddyGPT/usecases/web_api_testing/documentation_handler.py

@@ -61,7 +61,7 @@ def update_openapi_spec(self, resp, result):
                 if path not in self.openapi_spec['endpoints']:
                     self.openapi_spec['endpoints'][path] = {}
                 # Update the method description within the path
-                example, reference = self.response_handler.parse_http_response_to_openapi_example(result, path, method)
+                example, reference, self.openapi_spec = self.response_handler.parse_http_response_to_openapi_example(self.openapi_spec, result, path, method)
                 if example != None or reference != None:
                     self.openapi_spec['endpoints'][path][method.lower()] = {
                     "summary": f"{method} operation on {path}",
@@ -122,34 +122,3 @@ def check_openapi_spec(self, note):
         yaml_file_assistant = YamlFileAssistant(self.file_path, self.llm_handler)
         yaml_file_assistant.run(description)
 
-
-    def get_response_for_prompt(self, prompt):
-        """
-        Sends a prompt to OpenAI's API and retrieves the response.
-
-        Args:
-            prompt (str): The prompt to be sent to the API.
-
-        Returns:
-            str: The response from the API.
-        """
-        messages = [{"role": "user",
-                     "content": [{"type": "text", "text": prompt}
-                                 ]
-                     }
-                    ]
-
-
-        tic = time.perf_counter()
-        #print(f'shorten prompt: {prompt}')
-        response, completion = self.llm_handler.call_llm(messages)
-        toc = time.perf_counter()
-        # Update history
-        message = completion.choices[0].message
-        #print(f'Message: {message}')
-        command = pydantic_core.to_json(response).decode()
-        #print(f'response:{response}')
-        response_text = response.execute()
-        #print(f'[Response]: {response_text[:20]}')
-
-        return response_text

src/hackingBuddyGPT/usecases/web_api_testing/response_handler.py

@@ -4,10 +4,31 @@
 
 
 class ResponseHandler(object):
-    def __init__(self, name):
-        self.name = name
+    def __init__(self, llm_handler):
+        self.llm_handler = llm_handler
 
+    def get_response_for_prompt(self, prompt):
+        """
+        Sends a prompt to OpenAI's API and retrieves the response.
+
+        Args:
+            prompt (str): The prompt to be sent to the API.
+
+        Returns:
+            str: The response from the API.
+        """
+        messages = [{"role": "user",
+                     "content": [{"type": "text", "text": prompt}
+                                 ]
+                     }
+                    ]
+
+
+        response, completion = self.llm_handler.call_llm(messages)
+
+        response_text = response.execute()
 
+        return response_text
     def parse_http_status_line(self, status_line):
         if status_line == "Not a valid HTTP method":
             return status_line
@@ -37,7 +58,7 @@ def extract_response_example(self, html_content):
         # Format the response example
         return json.loads(result_text)
 
-    def parse_http_response_to_openapi_example(self, http_response, path, method):
+    def parse_http_response_to_openapi_example(self, openapi_spec, http_response, path, method):
         if method == "DELETE" or method == "PUT":
             print(f'http response: {http_response}')
         # Extract headers and body from the HTTP response
@@ -49,7 +70,7 @@ def parse_http_response_to_openapi_example(self, http_response, path, method):
             body_dict = json.loads(body)
         except json.decoder.JSONDecodeError:
             return None, None
-        reference, object_name = self.parse_http_response_to_schema(body_dict, path)
+        reference, object_name, openapi_spec = self.parse_http_response_to_schema(openapi_spec, body_dict, path)
 
         entry_dict = {}
         # Build the OpenAPI response example
@@ -63,11 +84,11 @@ def parse_http_response_to_openapi_example(self, http_response, path, method):
                 entry_dict[key] = {"value": entry}
                 self.llm_handler.add_created_object(entry_dict[key], object_name)
 
-        return entry_dict, reference
+        return entry_dict, reference, openapi_spec
     def extract_description(self, note):
         return note.action.content
 
-    def parse_http_response_to_schema(self, body_dict, path):
+    def parse_http_response_to_schema(self, openapi_spec ,body_dict, path):
         # Create object name
         object_name = path.split("/")[1].capitalize()
         object_name = object_name[:len(object_name) - 1]
@@ -87,14 +108,14 @@ def parse_http_response_to_schema(self, body_dict, path):
 
         object_dict = {"type": "object", "properties": properties_dict}
 
-        if not object_name in self.openapi_spec["components"]["schemas"].keys():
-            self.openapi_spec["components"]["schemas"][object_name] = object_dict
+        if not object_name in openapi_spec["components"]["schemas"].keys():
+            openapi_spec["components"]["schemas"][object_name] = object_dict
 
-        schemas = self.openapi_spec["components"]["schemas"]
+        schemas = openapi_spec["components"]["schemas"]
         self.schemas = schemas
         print(f'schemas: {schemas}')
         reference = "#/components/schemas/" + object_name
-        return reference, object_name
+        return reference, object_name, openapi_spec
     def read_yaml_to_string(self, filepath):
         """
         Reads a YAML file and returns its contents as a string.

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -53,8 +53,8 @@ class SimpleWebAPIDocumentation(RoundBasedUseCase):
     def init(self):
         super().init()
         self._setup_capabilities()
-        self.response_handler = ResponseHandler("doc")
         self.llm_handler = LLMHandler(self.llm, self._capabilities)
+        self.response_handler = ResponseHandler(self.llm_handler)
         self._setup_initial_prompt()
         self.documentation_handler = DocumentationHandler(self.llm_handler, self.response_handler)
 

src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py

@@ -1,9 +1,10 @@
 import time
-
 from dataclasses import dataclass, field
+from typing import List, Any, Union, Dict
+
+import pydantic_core
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionMessage
 from rich.panel import Panel
-from typing import List, Any, Union, Dict
 
 from hackingBuddyGPT.capabilities import Capability
 from hackingBuddyGPT.capabilities.capability import capabilities_to_action_model
@@ -14,35 +15,27 @@
 from hackingBuddyGPT.usecases.web_api_testing.llm_handler import LLMHandler
 from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy
 from hackingBuddyGPT.usecases.web_api_testing.response_handler import ResponseHandler
-from hackingBuddyGPT.utils import LLMResult, tool_message, ui
+from hackingBuddyGPT.utils import tool_message
 from hackingBuddyGPT.utils.configurable import parameter
 from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 from hackingBuddyGPT.usecases.base import use_case
 
-import pydantic_core
-
 Prompt = List[Union[ChatCompletionMessage, ChatCompletionMessageParam]]
 Context = Any
 
-
-@use_case("simple_web_api_testing", "Minimal implementation of a web api testing use case")
+@use_case("simple_web_api_testing", "Minimal implementation of a web API testing use case")
 @dataclass
 class SimpleWebAPITesting(RoundBasedUseCase):
     llm: OpenAILib
     host: str = parameter(desc="The host to test", default="https://jsonplaceholder.typicode.com")
-    # Parameter specifying the pattern description for expected HTTP methods in the API response
     http_method_description: str = parameter(
         desc="Pattern description for expected HTTP methods in the API response",
         default="A string that represents an HTTP method (e.g., 'GET', 'POST', etc.)."
     )
-
-    # Parameter specifying the template used to format HTTP methods in API requests
     http_method_template: str = parameter(
         desc="Template used to format HTTP methods in API requests. The {method} placeholder will be replaced by actual HTTP method names.",
         default="{method}"
     )
-
-    # Parameter specifying the expected HTTP methods as a comma-separated list
     http_methods: str = parameter(
         desc="Comma-separated list of HTTP methods expected to be used in the API response.",
         default="GET,POST,PUT,DELETE"
@@ -54,44 +47,82 @@ class SimpleWebAPITesting(RoundBasedUseCase):
     _all_http_methods_found: bool = False
 
     def init(self):
+        """
+        Initializes the SimpleWebAPITesting use case by setting up the context, response handler,
+        LLM handler, capabilities, and the initial prompt.
+        """
         super().init()
         self._context["host"] = self.host
-        self._setup_capabilities()
-        self.response_handler = ResponseHandler("doc")
         self.llm_handler = LLMHandler(self.llm, self._capabilities)
+        self.response_handler = ResponseHandler(self.llm_handler)
+        self._setup_capabilities()
         self._setup_initial_prompt()
-        self.response_handler = ResponseHandler("test")
 
     def _setup_initial_prompt(self):
+        """
+        Sets up the initial prompt for the LLM. The prompt provides instructions for the LLM
+        to evaluate the REST APIs of the specified host and to document findings.
+        """
         initial_prompt = {
             "role": "system",
-            "content": f"You're a seasoned API tester tasked with evaluating the REST APIs of a website hosted at {self.host}. "
-                       f"Your primary objective is to thoroughly explore the APIs, understanding their endpoints, parameters, and responses, and then devise and execute potential exploits.\n"
-                       f"Be diligent in documenting your findings as you navigate through the APIs..\n"
-                       f"Avoid resorting to brute-force tactics. All necessary information can be obtained from the API endpoints themselves.\n"
-                       f"Remember, if you encounter a http methods ({self.http_method_description}), promptly submit it as it is of utmost importance."
+            "content": (
+                f"You're a seasoned API tester tasked with evaluating the REST APIs of a website hosted at {self.host}. "
+                f"Your primary objective is to thoroughly explore the APIs, understanding their endpoints, parameters, and responses, "
+                f"and then devise and execute potential exploits. Be diligent in documenting your findings as you navigate through the APIs. "
+                f"Avoid resorting to brute-force tactics. All necessary information can be obtained from the API endpoints themselves. "
+                f"Remember, if you encounter an HTTP method ({self.http_method_description}), promptly submit it as it is of utmost importance."
+            )
         }
         self._prompt_history.append(initial_prompt)
-        self.prompt_engineer = PromptEngineer(strategy=PromptStrategy.CHAIN_OF_THOUGHT, llm_handler=self.llm_handler,
-                                              history=self._prompt_history, schemas={}, response_handler= self.response_handler)
+        self.prompt_engineer = PromptEngineer(
+            strategy=PromptStrategy.CHAIN_OF_THOUGHT, llm_handler=self.llm_handler,
+            history=self._prompt_history, schemas={}, response_handler=self.response_handler
+        )
 
     def all_http_methods_found(self):
+        """
+        Handles the event when all HTTP methods are found. Displays a congratulatory message
+        and sets the _all_http_methods_found flag to True.
+        """
         self.console.print(Panel("All HTTP methods found! Congratulations!", title="system"))
         self._all_http_methods_found = True
+
     def _setup_capabilities(self):
-        sett = {self.http_method_template.format(method=method) for method in self.http_methods.split(",")}
+        """
+        Sets up the capabilities required for the use case. Initializes HTTP request capabilities,
+        note recording capabilities, and HTTP method submission capabilities based on the provided
+        configuration.
+        """
+        methods_set = {self.http_method_template.format(method=method) for method in self.http_methods.split(",")}
         notes = self._context["notes"]
         self._capabilities = {
-            "submit_http_method": SubmitHTTPMethod(self.http_method_description, sett),
+            "submit_http_method": SubmitHTTPMethod(self.http_method_description, methods_set),
             "http_request": HTTPRequest(self.host),
             "record_note": RecordNote(notes)
         }
 
     def perform_round(self, turn: int, FINAL_ROUND=30):
+        """
+        Performs a single round of interaction with the LLM. Generates a prompt, sends it to the LLM,
+        and handles the response.
+
+        Args:
+            turn (int): The current round number.
+            FINAL_ROUND (int, optional): The final round number. Defaults to 30.
+        """
         prompt = self.prompt_engineer.generate_prompt(doc=True)
         response, completion = self.llm_handler.call_llm(prompt)
         self._handle_response(completion, response)
+
     def _handle_response(self, completion, response):
+        """
+        Handles the response from the LLM. Parses the response, executes the necessary actions,
+        and updates the prompt history.
+
+        Args:
+            completion (Any): The completion object from the LLM.
+            response (Any): The response object from the LLM.
+        """
         message = completion.choices[0].message
         tool_call_id = message.tool_calls[0].id
         command = pydantic_core.to_json(response).decode()
@@ -103,8 +134,5 @@ def _handle_response(self, completion, response):
             self.console.print(Panel(result[:30], title="tool"))
             result_str = self.response_handler.parse_http_status_line(result)
             self._prompt_history.append(tool_message(result_str, tool_call_id))
-            invalid_flags = ["recorded","Not a valid HTTP method" ]
-        return self._all_http_methods_found
-
-
 
+        return self._all_http_methods_found