Adjusted to only record valid information of rest api

src/hackingBuddyGPT/capabilities/python_test_case.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass, field
+from typing import Dict, Any, List, Tuple
+from hackingBuddyGPT.capabilities import Capability
+
+
+@dataclass
+class PythonTestCase(Capability):
+    description: str
+    input: Dict[str, Any] = field(default_factory=dict)
+    expected_output: Dict[str, Any] = field(default_factory=dict)
+    registry: List[Tuple[str, str]] = field(default_factory=list)
+
+    def describe(self) -> str:
+        """
+        Returns a description of the test case.
+        """
+        return f"Test Case: {self.description}\nInput: {self.input}\nExpected Output: {self.expected_output}"
+    def __call__(self, title: str, content: str) -> str:
+        self.registry.append((title, content))
+        return f" Test Case:\n{title}: {content}"

src/hackingBuddyGPT/usecases/web_api_testing/documentation/openapi_specification_handler.py

@@ -81,7 +81,7 @@ def update_openapi_spec(self, resp, result, result_str):
             path = request.path
             method = request.method
 
-            if not path or not method :
+            if not path or not method or path == "/":
                 return list(self.openapi_spec["endpoints"].keys())
             if "1" in path:
                 path = path.replace("1", ":id")
@@ -93,9 +93,11 @@ def update_openapi_spec(self, resp, result, result_str):
             main_path = path if len(path_parts) > 1 else ""
 
             # Initialize the path if it's not present and is valid
-            if path not in endpoints and main_path:
+            if path not in endpoints and main_path and status_code == "200":
                 endpoints[path] = {}
                 endpoint_methods[path] = []
+            if path not in endpoints:
+                return list(self.openapi_spec["endpoints"].keys())
 
             # Parse the response into OpenAPI example and reference
             example, reference, self.openapi_spec = self.response_handler.parse_http_response_to_openapi_example(

src/hackingBuddyGPT/usecases/web_api_testing/prompt_generation/prompt_generation_helper.py

@@ -142,9 +142,10 @@ def _get_initial_documentation_steps(self, common_steps, strategy):
         Returns:
             list: A list of initial steps combined with common steps.
         """
+        endpoints = list(set([ endpoint.replace(":id", "1") for endpoint in self.found_endpoints] + ['/']))
         documentation_steps = [
             f"""Identify all available endpoints via GET Requests. 
-            Exclude those in this list: {[ endpoint.replace(":id", "1") for endpoint in self.found_endpoints]} 
+            Exclude those in this list: {endpoints} 
             and endpoints that match this pattern: '/resource/number' where 'number' is greater than 1 (e.g., '/todos/2', '/todos/3').
             Only include endpoints where the number is 1 or the endpoint does not end with a number at all.
 

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -36,7 +36,7 @@ class SimpleWebAPIDocumentation(Agent):
     """
 
     llm: OpenAILib
-    host: str = parameter(desc="The host to test", default="https://jsonplaceholder.typicode.com")
+    host: str = parameter(desc="The host to test", default="https://84e9-213-255-219-62.ngrok-free.app")
     _prompt_history: Prompt = field(default_factory=list)
     _context: Context = field(default_factory=lambda: {"notes": list()})
     _capabilities: Dict[str, Capability] = field(default_factory=dict)

src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py

@@ -1,31 +1,33 @@
 import os.path
 from dataclasses import field
+from datetime import datetime
 from typing import Any, Dict, List
 
 import pydantic_core
 from rich.panel import Panel
 
 from hackingBuddyGPT.capabilities import Capability
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
+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.web_api_testing.prompt_generation.information.prompt_information import PromptContext, \
     PromptPurpose
-from hackingBuddyGPT.usecases.web_api_testing.utils.custom_datatypes import Prompt, Context
 from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing import OpenAPISpecificationParser
 from hackingBuddyGPT.usecases.web_api_testing.documentation.report_handler import ReportHandler
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptContext
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompt_engineer import PromptEngineer, PromptStrategy
 from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_handler import ResponseHandler
+from hackingBuddyGPT.usecases.web_api_testing.testing.test_handler import TestHandler
 from hackingBuddyGPT.usecases.web_api_testing.utils.custom_datatypes import Context, Prompt
 from hackingBuddyGPT.usecases.web_api_testing.utils.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
-openapi_spec_filename = "src/hackingBuddyGPT/usecases/web_api_testing/documentation/openapi_spec/openapi_spec_2024-09-03_10-22-09.yaml"
+openapi_spec_filename = "src/hackingBuddyGPT/usecases/web_api_testing/documentation/openapi_spec/in_context/openapi_spec_2024-10-16_15-36-11.yaml"
 
 
 class SimpleWebAPITesting(Agent):
@@ -60,7 +62,7 @@ class SimpleWebAPITesting(Agent):
     )
 
     _prompt_history: Prompt = field(default_factory=list)
-    _context: Context = field(default_factory=lambda: {"notes": list()})
+    _context: Context = field(default_factory=lambda: {"notes": list(), "test_cases":list})
     _capabilities: Dict[str, Capability] = field(default_factory=dict)
     _all_http_methods_found: bool = False
 
@@ -77,9 +79,12 @@ def init(self) -> None:
         self._llm_handler: LLMHandler = LLMHandler(self.llm, self._capabilities)
         self._response_handler: ResponseHandler = ResponseHandler(self._llm_handler)
         self._report_handler: ReportHandler = ReportHandler()
+        self._test_handler: TestHandler = TestHandler(self._llm_handler)
         self._setup_initial_prompt()
         self.purpose =  PromptPurpose.AUTHENTICATION_AUTHORIZATION
 
+
+
     def _setup_initial_prompt(self) -> None:
         """
         Sets up the initial prompt for the LLM. The prompt provides instructions for the LLM
@@ -106,9 +111,9 @@ def _setup_initial_prompt(self) -> None:
             history=self._prompt_history,
             handlers=handlers,
             context=PromptContext.PENTESTING,
-            rest_api=self.host,
             schemas=schemas,
-            endpoints= endpoints
+            endpoints= endpoints,
+
         )
 
     def all_http_methods_found(self) -> None:
@@ -129,10 +134,12 @@ def _setup_capabilities(self) -> None:
             self.http_method_template.format(method=method) for method in self.http_methods.split(",")
         }
         notes: List[str] = self._context["notes"]
+        test_cases = self._context["test_cases"]
         self._capabilities = {
             "submit_http_method": HTTPRequest(self.host),
             "http_request": HTTPRequest(self.host),
             "record_note": RecordNote(notes),
+            "test_cases": PythonTestCase(test_cases)
         }
 
     def perform_round(self, turn: int) -> None:
@@ -186,6 +193,8 @@ def _handle_response(self, completion: Any, response: Any, purpose: str) -> None
             self._prompt_history.append(tool_message(self._response_handler.extract_key_elements_of_response(result), tool_call_id))
 
             analysis = self._response_handler.evaluate_result(result=result, prompt_history=self._prompt_history)
+
+            self._test_handler.generate_and_save_test_cases(analysis=analysis, endpoint=response.action.path, method=response.action.method, prompt_history= self._prompt_history)
             self._report_handler.write_analysis_to_report(analysis=analysis, purpose=self.prompt_engineer.purpose)
 
         self.all_http_methods_found()

src/hackingBuddyGPT/usecases/web_api_testing/testing/test_handler.py

@@ -0,0 +1,192 @@
+import json
+import os
+import re
+from datetime import datetime
+from typing import Any, Dict, Tuple
+
+import pydantic_core
+
+
+class TestHandler(object):
+
+    def __init__(self, llm_handler):
+        self._llm_handler = llm_handler
+        current_path = os.path.dirname(os.path.abspath(__file__))
+        self.test_path = os.path.join(current_path, "tests")
+        self.filename = f"test{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
+
+        self.file = os.path.join(self.test_path, self.filename)
+
+    def parse_test_case(self, note: str) -> Dict[str, Any]:
+            """
+            Parses a note containing a test case into a structured format.
+
+            Args:
+                note (str): The note string containing the test case information.
+
+            Returns:
+                Dict[str, Any]: The parsed test case in a structured format.
+            """
+            # Regular expressions to extract the method, endpoint, input, and expected output
+            method_endpoint_pattern = re.compile(r"Test Case for (\w+) (\/\S+):")
+            description_pattern = re.compile(r"Description: (.+)")
+            input_data_pattern = re.compile(r"Input Data: (\{.*\})")
+            expected_output_pattern = re.compile(r"Expected Output: (.+)")
+
+            # Extract method and endpoint
+            method_endpoint_match = method_endpoint_pattern.search(note)
+            if method_endpoint_match:
+                method, endpoint = method_endpoint_match.groups()
+            else:
+                raise ValueError("Method and endpoint not found in the note")
+
+            # Extract description
+            description_match = description_pattern.search(note)
+            description = description_match.group(1) if description_match else "No description found"
+
+            # Extract input data
+            input_data_match = input_data_pattern.search(note)
+            input_data = input_data_match.group(1) if input_data_match else "{}"
+
+            # Extract expected output
+            expected_output_match = expected_output_pattern.search(note)
+            expected_output = expected_output_match.group(1) if expected_output_match else "No expected output found"
+
+            # Construct the structured test case
+            test_case = {
+                "description": f"Test case for {method} {endpoint}",
+                "input": input_data,
+                "expected_output": expected_output
+            }
+
+            return test_case
+
+    def generate_test_case(self, analysis: str, endpoint: str, method: str, prompt_history) -> Tuple[str, Dict[str, Any]]:
+        """
+        Generates a test case based on the provided analysis of the API response.
+
+        Args:
+            analysis (str): Analysis of the API response and its behavior.
+            endpoint (str): The API endpoint being tested.
+            method (str): The HTTP method to use in the test case.
+
+        Returns:
+            Tuple[str, Dict[str, Any]]: A description of the test case and the payload.
+        """
+        prompt_text = f"""
+           Based on the following analysis of the API response, generate a detailed test case:
+
+           Analysis: {analysis}
+
+           Endpoint: {endpoint}
+           HTTP Method: {method}
+
+           The test case should include:
+           - Description of the test.
+           - Example input data in JSON format.
+           - Expected result or assertion.
+
+           Example Format:
+           {{
+               "description": "Test case for {method} {endpoint}",
+               "input": {{}},
+               "expected_output": {{}}
+           }}
+           """
+        prompt_history.append({"role": "system", "content": prompt_text})
+
+        response, completion = self._llm_handler.call_llm(prompt_history)
+        message = completion.choices[0].message
+        tool_call_id: str = message.tool_calls[0].id
+        command: str = pydantic_core.to_json(response).decode()
+        result: Any = response.execute()
+        test_case = self.parse_test_case(result)
+        # Extract the structured test case if possible
+        try:
+            test_case_dict = json.loads(test_case)
+        except json.JSONDecodeError:
+            raise ValueError("LLM-generated test case is not valid JSON")
+
+        return test_case_dict["description"], test_case_dict
+
+    def write_test_case_to_file(self, description: str, test_case: Dict[str, Any]) -> None:
+        """
+        Writes a generated test case to a specified file.
+
+        Args:
+            description (str): Description of the test case.
+            test_case (Dict[str, Any]): The test case including input and expected output.
+            output_file (str): The file path where the test case should be saved.
+        """
+        test_case_entry = {
+            "description": description,
+            "test_case": test_case
+        }
+
+        with open(self.file + ".json", "a") as f:
+            f.write(json.dumps(test_case_entry, indent=2) + "\n\n")
+
+        print((f"Test case written to {self.file}"))
+
+    def write_pytest_case(self, description: str, test_case: Dict[str, Any]) -> None:
+        """
+        Writes a pytest-compatible test case to a Python file using LLM for code generation.
+
+        Args:
+            description (str): Description of the test case.
+            test_case (Dict[str, Any]): The test case including input and expected output.
+        """
+        # Construct a prompt to guide the LLM in generating the test code.
+        prompt = f"""
+        You are an expert Python developer specializing in writing automated tests using pytest.
+        Based on the following details, generate a pytest-compatible test function:
+
+        Description: {description}
+
+        Test Case:
+        - Endpoint: {test_case['endpoint']}
+        - HTTP Method: {test_case['method'].upper()}
+        - Input Data: {json.dumps(test_case.get("input", {}), indent=4)}
+        - Expected Status Code: {test_case['expected_output'].get('status_code', 200)}
+        - Expected Response Body: {json.dumps(test_case['expected_output'].get('body', {}), indent=4)}
+
+        The generated test function should:
+        - Use the 'requests' library to make the HTTP request.
+        - Include assertions for the status code and the response body.
+        - Be properly formatted and ready to use with pytest.
+        - Include a docstring with the test description.
+
+        Example Format:
+        ```
+        import requests
+        import pytest
+
+        @pytest.mark.api
+        def test_example():
+            \"\"\"Description of the test.\"\"\"
+            # Test implementation here
+        ```
+        """
+
+        # Call the LLM to generate the test function.
+        response = self._llm_handler.call_llm(prompt)
+        test_function = response['choices'][0]['text']
+
+        # Write the generated test function to a Python file.
+        with open(self.file + ".py", "a") as f:
+            f.write(test_function)
+
+        print(f"Pytest case written to {self.file}.py")
+
+    def generate_and_save_test_cases(self, analysis: str, endpoint: str, method: str, prompt_history) -> None:
+        """
+        Generates test cases based on the analysis and saves them as pytest-compatible tests.
+
+        Args:
+            analysis (str): Analysis of the API response.
+            endpoint (str): The endpoint being tested.
+            method (str): The HTTP method used for testing.
+        """
+        description, test_case = self.generate_test_case(analysis, endpoint, method, prompt_history)
+        self.write_test_case_to_file(description, test_case)
+        self.write_pytest_case(description, test_case)

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

@@ -37,7 +37,7 @@ def __init__(self, llm_handler, response_handler):
                 "version": "1.0",
                 "description": "Automatically generated description of the API."
             },
-            "servers": [{"url": "https://jsonplaceholder.typicode.com"}],
+            "servers": [{"url": "https://localhost:8080"}],
             "endpoints": {},
             "components": {"schemas": {}}
         }
@@ -123,6 +123,6 @@ def check_openapi_spec(self, note):
             note (object): The note object containing the description of the API.
         """
         description = self.response_handler.extract_description(note)
-        from hackingBuddyGPT.usecases.web_api_testing.utils.yaml_assistant import YamlFileAssistant
+        from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing.yaml_assistant import YamlFileAssistant
         yaml_file_assistant = YamlFileAssistant(self.file_path, self.llm_handler)
         yaml_file_assistant.run(description)

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

@@ -116,7 +116,8 @@ def get_created_objects(self) -> Dict[str, List[Any]]:
         return self.created_objects
 
     def adjust_prompt_based_on_token(self, prompt: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        prompt.reverse()
+        if not isinstance(prompt, str):
+            prompt.reverse()
         tokens = 0
         max_tokens = 10000
         for item in prompt:
@@ -131,7 +132,8 @@ def adjust_prompt_based_on_token(self, prompt: List[Dict[str, Any]]) -> List[Dic
                     continue
 
         print(f"tokens:{tokens}")
-        prompt.reverse()
+        if not isinstance(prompt, str):
+            prompt.reverse()
         return prompt
 
     def get_num_tokens(self, content: str) -> int: