Merge pull request #82 from ipa-lab/web_api_testing

Improved WebAPITesting

### New Features:
- **ResponseAnalyzerWithLLM**: Introduced a new `ResponseAnalyzerWithLLM` to replace the previous `ResponseAnalyzer`, offering enhanced capabilities for large language model-based analysis.
- **Task and State Planning**: Divided prompts into distinct tasks and state planning for improved processing.

### Improvements:
- **Code Optimization**: Multiple optimizations were performed to improve code efficiency and performance.
- **Refactoring**: Significant restructuring was done to improve code reusability and maintainability, particularly in the `web_api_testing` and `prompt_helper` components.
- **Improved Documentation**: Expanded and clarified documentation throughout the codebase, including updates to `web_api_documentation` and additional data types.
- **Enhanced Testing**: Added and adjusted tests for the new features and refactored components to ensure robust testing coverage.
- **Pentesting Improvements**: Integrated a more sophisticated response analyzer for better pentesting capabilities and optimized token counting.

### Bug Fixes:
- **Test Fixes**: Resolved issues with tests after code restructures and ensured all tests are now passing post-optimization.
- **.gitignore Adjustments**: Updated `.gitignore` to exclude unnecessary files from version control.

Author: andreashappe

.gitignore

@@ -13,4 +13,5 @@ dist/
 .coverage
 src/hackingBuddyGPT/usecases/web_api_testing/openapi_spec/
 src/hackingBuddyGPT/usecases/web_api_testing/converted_files/
-/src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_spec/
+/src/hackingBuddyGPT/usecases/web_api_testing/documentation/openapi_spec/
+/src/hackingBuddyGPT/usecases/web_api_testing/documentation/reports/

src/hackingBuddyGPT/cli/wintermute.py

@@ -8,13 +8,12 @@ def main():
     parser = argparse.ArgumentParser()
     subparser = parser.add_subparsers(required=True)
     for name, use_case in use_cases.items():
-        subb = subparser.add_parser(
+        use_case.build_parser(subparser.add_parser(
             name=use_case.name,
             help=use_case.description
-        )
-        use_case.build_parser(subb)
-    x= sys.argv[1:]
-    parsed = parser.parse_args(x)
+        ))
+
+    parsed = parser.parse_args(sys.argv[1:])
     instance = parsed.use_case(parsed)
     instance.init()
     instance.run()

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

@@ -0,0 +1,2 @@
+from .openapi_specification_handler import OpenAPISpecificationHandler
+from .report_handler import ReportHandler

src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_specification_manager.py renamed to src/hackingBuddyGPT/usecases/web_api_testing/documentation/openapi_specification_handler.py

@@ -2,8 +2,14 @@
 import yaml
 from datetime import datetime
 from hackingBuddyGPT.capabilities.yamlFile import YAMLFile
-
-class OpenAPISpecificationManager:
+from collections import defaultdict
+import pydantic_core
+from rich.panel import Panel
+
+from hackingBuddyGPT.usecases.web_api_testing.response_processing import ResponseHandler
+from hackingBuddyGPT.usecases.web_api_testing.utils import LLMHandler
+from hackingBuddyGPT.utils import tool_message
+class OpenAPISpecificationHandler(object):
     """
     Handles the generation and updating of an OpenAPI specification document based on dynamic API responses.
 
@@ -19,7 +25,7 @@ class OpenAPISpecificationManager:
         _capabilities (dict): A dictionary to store capabilities related to YAML file handling.
     """
 
-    def __init__(self, llm_handler, response_handler):
+    def __init__(self, llm_handler: LLMHandler, response_handler: ResponseHandler):
         """
         Initializes the handler with a template OpenAPI specification.
 
@@ -43,7 +49,6 @@ def __init__(self, llm_handler, response_handler):
             "components": {"schemas": {}}
         }
         self.llm_handler = llm_handler
-        #self.api_key = llm_handler.llm.api_key
         current_path = os.path.dirname(os.path.abspath(__file__))
         self.file_path = os.path.join(current_path, "openapi_spec")
         self.file = os.path.join(self.file_path, self.filename)
@@ -149,6 +154,49 @@ 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.utils.documentation.parsing.yaml_assistant import YamlFileAssistant
         yaml_file_assistant = YamlFileAssistant(self.file_path, self.llm_handler)
         yaml_file_assistant.run(description)
+
+
+    def _update_documentation(self, response, result, prompt_engineer):
+        prompt_engineer.prompt_helper.found_endpoints = self.update_openapi_spec(response,
+                                                                                                            result)
+        self.write_openapi_to_yaml()
+        prompt_engineer.prompt_helper.schemas = self.schemas
+
+        http_methods_dict = defaultdict(list)
+        for endpoint, methods in self.endpoint_methods.items():
+            for method in methods:
+                http_methods_dict[method].append(endpoint)
+
+        prompt_engineer.prompt_helper.endpoint_found_methods = http_methods_dict
+        prompt_engineer.prompt_helper.endpoint_methods = self.endpoint_methods
+        return prompt_engineer
+
+    def document_response(self, completion, response, log, prompt_history, prompt_engineer):
+            message = completion.choices[0].message
+            tool_call_id = message.tool_calls[0].id
+            command = pydantic_core.to_json(response).decode()
+
+            log.console.print(Panel(command, title="assistant"))
+            prompt_history.append(message)
+
+            with log.console.status("[bold green]Executing that command..."):
+                result = response.execute()
+                log.console.print(Panel(result[:30], title="tool"))
+                result_str = self.response_handler.parse_http_status_line(result)
+                prompt_history.append(tool_message(result_str, tool_call_id))
+
+                invalid_flags = {"recorded", "Not a valid HTTP method", "404", "Client Error: Not Found"}
+                if not result_str in invalid_flags or any(flag in result_str for flag in invalid_flags):
+                    prompt_engineer = self._update_documentation(response, result, prompt_engineer)
+
+            return log, prompt_history, prompt_engineer
+
+    def found_all_endpoints(self):
+        if len(self.endpoint_methods.items())< 10:
+            return False
+        else:
+            return True
+

src/hackingBuddyGPT/usecases/web_api_testing/documentation/parsing/__init__.py

@@ -0,0 +1,3 @@
+from .openapi_converter import OpenAPISpecificationConverter
+from .openapi_parser import OpenAPISpecificationParser
+from .yaml_assistant import YamlFileAssistant

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

@@ -1,87 +1,81 @@
 import yaml
+from typing import Dict, List, Union
 
 class OpenAPISpecificationParser:
     """
     OpenAPISpecificationParser is a class for parsing and extracting information from an OpenAPI specification file.
 
     Attributes:
         filepath (str): The path to the OpenAPI specification YAML file.
-        api_data (dict): The parsed data from the YAML file.
+        api_data (Dict[str, Union[Dict, List]]): The parsed data from the YAML file.
     """
 
-    def __init__(self, filepath):
+    def __init__(self, filepath: str):
         """
         Initializes the OpenAPISpecificationParser with the specified file path.
 
         Args:
             filepath (str): The path to the OpenAPI specification YAML file.
         """
-        self.filepath = filepath
-        self.api_data = self.load_yaml()
+        self.filepath: str = filepath
+        self.api_data: Dict[str, Union[Dict, List]] = self.load_yaml()
 
-    def load_yaml(self):
+    def load_yaml(self) -> Dict[str, Union[Dict, List]]:
         """
         Loads YAML data from the specified file.
 
         Returns:
-            dict: The parsed data from the YAML file.
+            Dict[str, Union[Dict, List]]: The parsed data from the YAML file.
         """
         with open(self.filepath, 'r') as file:
             return yaml.safe_load(file)
 
-    def get_servers(self):
+    def _get_servers(self) -> List[str]:
         """
         Retrieves the list of server URLs from the OpenAPI specification.
 
         Returns:
-            list: A list of server URLs.
+            List[str]: A list of server URLs.
         """
         return [server['url'] for server in self.api_data.get('servers', [])]
 
-    def get_paths(self):
+    def get_paths(self) -> Dict[str, Dict[str, Dict]]:
         """
         Retrieves all API paths and their methods from the OpenAPI specification.
 
         Returns:
-            dict: A dictionary with API paths as keys and methods as values.
+            Dict[str, Dict[str, Dict]]: A dictionary with API paths as keys and methods as values.
         """
-        paths_info = {}
-        paths = self.api_data.get('paths', {})
+        paths_info: Dict[str, Dict[str, Dict]] = {}
+        paths: Dict[str, Dict[str, Dict]] = self.api_data.get('paths', {})
         for path, methods in paths.items():
             paths_info[path] = {method: details for method, details in methods.items()}
         return paths_info
 
-    def get_operations(self, path):
+    def _get_operations(self, path: str) -> Dict[str, Dict]:
         """
         Retrieves operations for a specific path from the OpenAPI specification.
 
         Args:
             path (str): The API path to retrieve operations for.
 
         Returns:
-            dict: A dictionary with methods as keys and operation details as values.
+            Dict[str, Dict]: A dictionary with methods as keys and operation details as values.
         """
         return self.api_data['paths'].get(path, {})
 
-    def print_api_details(self):
+    def _print_api_details(self) -> None:
         """
         Prints details of the API extracted from the OpenAPI document, including title, version, servers,
         paths, and operations.
         """
         print("API Title:", self.api_data['info']['title'])
         print("API Version:", self.api_data['info']['version'])
-        print("Servers:", self.get_servers())
+        print("Servers:", self._get_servers())
         print("\nAvailable Paths and Operations:")
         for path, operations in self.get_paths().items():
             print(f"\nPath: {path}")
             for operation, details in operations.items():
                 print(f"  Operation: {operation.upper()}")
                 print(f"    Summary: {details.get('summary')}")
                 print(f"    Description: {details['responses']['200']['description']}")
-
-# Usage example
-if __name__ == '__main__':
-    openapi_parser = OpenAPISpecificationParser(
-        '/hackingBuddyGPT/usecases/web_api_testing/openapi_spec/openapi_spec_2024-06-13_17-16-25.yaml'
-    )
-    openapi_parser.print_api_details()

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

@@ -0,0 +1,91 @@
+from openai import OpenAI
+from typing import Any
+
+
+class YamlFileAssistant:
+    """
+    YamlFileAssistant is a class designed to interact with a YAML file using OpenAI's API.
+
+    Attributes:
+        yaml_file (str): The path to the YAML file that the assistant will analyze.
+        client (OpenAI): The OpenAI client used to interact with the OpenAI API.
+    """
+
+    def __init__(self, yaml_file: str, client: OpenAI):
+        """
+        Initializes the YamlFileAssistant with a specified YAML file and OpenAI client.
+
+        Args:
+            yaml_file (str): The path to the YAML file to be analyzed.
+            client (OpenAI): The OpenAI client used to interact with the OpenAI API.
+        """
+        self.yaml_file: str = yaml_file
+        self.client: OpenAI = client
+
+    def run(self, recorded_note: str) -> None:
+        """
+        Runs the assistant to analyze the YAML file based on a recorded note.
+
+        This method would typically interact with OpenAI's API to create an assistant,
+        upload the YAML file, analyze its contents, and generate responses. However, the
+        actual implementation is currently commented out.
+
+        Args:
+            recorded_note (str): A string containing the note or instructions for analysis.
+
+        Note:
+            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 "
+                         f"if the following information is contained in the provided yaml file. Information: {recorded_note}",
+            model="gpt-4o",
+            tools=[{"type": "file_search"}],
+        )
+
+        # Create a vector store called "Financial Statements"
+        vector_store = self.client.beta.vector_stores.create(name="Financial Statements")
+
+        # Ready the files for upload to OpenAI
+        file_streams = [open(self.yaml_file, "rb")]
+
+        # Use the upload and poll SDK helper to upload the files, add them to the vector store,
+        # and poll the status of the file batch for completion.
+        file_batch = self.client.beta.vector_stores.file_batches.upload_and_poll(
+            vector_store_id=vector_store.id, files=file_streams
+        )
+
+        # You can print the status and the file counts of the batch to see the result of this operation.
+        print(file_batch.status)
+        print(file_batch.file_counts)
+
+        assistant = self.client.beta.assistants.update(
+            assistant_id=assistant.id,
+            tool_resources={"file_search": {"vector_store_ids": [vector_store.id]}},
+        )
+
+        # Upload the user-provided file to OpenAI
+        message_file = self.client.files.create(
+            file=open("edgar/aapl-10k.pdf", "rb"), purpose="assistants"
+        )
+
+        # Create a thread and attach the file to the message
+        thread = self.client.beta.threads.create(
+            messages=[
+                {
+                    "role": "user",
+                    "content": "How many shares of AAPL were outstanding at the end of October 2023?",
+                    # Attach the new file to the message.
+                    "attachments": [
+                        {"file_id": message_file.id, "tools": [{"type": "file_search"}]}
+                    ],
+                }
+            ]
+        )
+
+        # The thread now has a vector store with that file in its tool resources.
+        print(thread.tool_resources.file_search)
+        '''

src/hackingBuddyGPT/usecases/web_api_testing/documentation/parsing/yaml_assistant.py

@@ -0,0 +1,62 @@
+import os
+from datetime import datetime
+import uuid
+from typing import List
+from enum import Enum
+
+class ReportHandler:
+    """
+    A handler for creating and managing report files that document operations and data.
+
+    Attributes:
+        file_path (str): The path to the directory where report files are stored.
+        report_name (str): The full path to the current report file being written to.
+        report (file): The file object for the report, opened for writing data.
+    """
+
+    def __init__(self):
+        """
+        Initializes the ReportHandler by setting up the file path for reports,
+        creating the directory if it does not exist, and preparing a new report file.
+        """
+        current_path: str = os.path.dirname(os.path.abspath(__file__))
+        self.file_path: str = os.path.join(current_path, "reports")
+
+        if not os.path.exists(self.file_path):
+            os.mkdir(self.file_path)
+
+        self.report_name: str = os.path.join(self.file_path, f"report_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.txt")
+        try:
+            self.report = open(self.report_name, "x")
+        except FileExistsError:
+            # Retry with a different name using a UUID to ensure uniqueness
+            self.report_name = os.path.join(self.file_path,
+                                            f"report_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{uuid.uuid4().hex}.txt")
+            self.report = open(self.report_name, "x")
+
+    def write_endpoint_to_report(self, endpoint: str) -> None:
+        """
+        Writes an endpoint string to the report file.
+
+        Args:
+            endpoint (str): The endpoint information to be recorded in the report.
+        """
+        with open(self.report_name, 'a') as report:
+            report.write(f'{endpoint}\n')
+
+    def write_analysis_to_report(self, analysis: List[str], purpose: Enum) -> None:
+        """
+        Writes an analysis result and its purpose to the report file.
+
+        Args:
+            analysis (List[str]): The analysis data to be recorded.
+            purpose (Enum): An enumeration that describes the purpose of the analysis.
+        """
+        with open(self.report_name, 'a') as report:
+            report.write(f'{purpose.name}:\n')
+            for item in analysis:
+                for line in item.split("\n"):
+                    if "note recorded" in line:
+                        continue
+                    else:
+                        report.write(line + "\n")