Merge branch 'development' of github.com:ipa-lab/hackingBuddyGPT into development

Author: andreashappe

.github/workflows/python-app.yml

@@ -5,9 +5,9 @@ name: Python application
 
 on:
   push:
-    branches: [ "main" ]
+    branches: [ "main", "development" ]
   pull_request:
-    branches: [ "main" ]
+    branches: [ "main", "development" ]
 
 permissions:
   contents: read

README.md

@@ -12,6 +12,7 @@ If you want to use hackingBuddyGPT and need help selecting the best LLM for your
 
 ## hackingBuddyGPT in the News
 
+- **upcoming** 2024-11-20: [Manuel Reinsperger](https://www.github.com/neverbolt) will present hackingBuddyGPT at the [European Symposium on Security and Artificial Intelligence (ESSAI)](https://essai-conference.eu/) 
 - 2024-07-26: The [GitHub Accelerator Showcase](https://github.blog/open-source/maintainers/github-accelerator-showcase-celebrating-our-second-cohort-and-whats-next/) features hackingBuddyGPT
 - 2024-07-24: [Juergen](https://github.com/citostyle) speaks at [Open Source + mezcal night @ GitHub HQ](https://lu.ma/bx120myg)
 - 2024-05-23: hackingBuddyGPT is part of [GitHub Accelerator 2024](https://github.blog/news-insights/company-news/2024-github-accelerator-meet-the-11-projects-shaping-open-source-ai/)
@@ -178,6 +179,9 @@ wintermute.py: error: the following arguments are required: {linux_privesc,windo
 
 # start wintermute, i.e., attack the configured virtual machine
 $ python wintermute.py minimal_linux_privesc
+
+# install dependencies for testing if you want to run the tests
+$ pip install .[testing]
 ~~~
 
 ## Publications about hackingBuddyGPT

pyproject.toml

@@ -57,6 +57,11 @@ pythonpath = "src"
 addopts = [
 	"--import-mode=importlib",
 ]
+[project.optional-dependencies]
+testing = [
+    'pytest',
+    'pytest-mock'
+]
 
 [project.scripts]
 wintermute = "hackingBuddyGPT.cli.wintermute:main"

src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py

@@ -1,6 +1,5 @@
 import nltk
 from nltk.tokenize import word_tokenize
-from nltk.corpus import stopwords
 from instructor.retry import InstructorRetryException
 
 
@@ -37,13 +36,11 @@ def __init__(self, strategy, llm_handler, history, schemas, response_handler):
         self.found_endpoints = ["/"]
         self.endpoint_methods = {}
         self.endpoint_found_methods = {}
-        model_name = "en_core_web_sm"
-
         # Check if the models are already installed
         nltk.download('punkt')
         nltk.download('stopwords')
         self._prompt_history = history
-        self.prompt = self._prompt_history
+        self.prompt = {self.round: {"content": "initial_prompt"}}
         self.previous_prompt = self._prompt_history[self.round]["content"]
         self.schemas = schemas
 
@@ -77,10 +74,6 @@ def generate_prompt(self, doc=False):
                     self.round = self.round +1
                     return self._prompt_history
 
-
-
-
-
     def in_context_learning(self, doc=False, hint=""):
         """
         Generates a prompt for in-context learning.
@@ -91,7 +84,14 @@ def in_context_learning(self, doc=False, hint=""):
         Returns:
             str: The generated prompt.
         """
-        return str("\n".join(self._prompt_history[self.round]["content"] + [self.prompt]))
+        history_content = [entry["content"] for entry in self._prompt_history]
+        prompt_content = self.prompt.get(self.round, {}).get("content", "")
+
+        # Add hint if provided
+        if hint:
+            prompt_content += f"\n{hint}"
+
+        return "\n".join(history_content + [prompt_content])
 
     def get_http_action_template(self, method):
         """Helper to construct a consistent HTTP action description."""
@@ -103,13 +103,45 @@ def get_http_action_template(self, method):
         else:
             return (
                 f"Create HTTPRequests of type {method} considering only the object with id=1 for the endpoint and understand the responses. Ensure that they are correct requests.")
-
-
+    def get_initial_steps(self, common_steps):
+            return [
+                "Identify all available endpoints via GET Requests. Exclude those in this list: {self.found_endpoints}",
+                "Note down the response structures, status codes, and headers for each endpoint.",
+                "For each endpoint, document the following details: URL, HTTP method, query parameters and path variables, expected request body structure for requests, response structure for successful and error responses."
+            ] + common_steps
+
+    def get_phase_steps(self, phase, common_steps):
+            if phase != "DELETE":
+                return [
+                    f"Identify for all endpoints {self.found_endpoints} excluding {self.endpoint_found_methods[phase]} a valid HTTP method {phase} call.",
+                    self.get_http_action_template(phase)
+                ] + common_steps
+            else:
+                return [
+                    "Check for all endpoints the DELETE method. Delete the first instance for all endpoints.",
+                    self.get_http_action_template(phase)
+                ] + common_steps
+
+    def get_endpoints_needing_help(self):
+            endpoints_needing_help = []
+            endpoints_and_needed_methods = {}
+            http_methods_set = {"GET", "POST", "PUT", "DELETE"}
+
+            for endpoint, methods in self.endpoint_methods.items():
+                missing_methods = http_methods_set - set(methods)
+                if len(methods) < 4:
+                    endpoints_needing_help.append(endpoint)
+                    endpoints_and_needed_methods[endpoint] = list(missing_methods)
+
+            if endpoints_needing_help:
+                first_endpoint = endpoints_needing_help[0]
+                needed_method = endpoints_and_needed_methods[first_endpoint][0]
+                return [
+                    f"For endpoint {first_endpoint} find this missing method: {needed_method}. If all the HTTP methods have already been found for an endpoint, then do not include this endpoint in your search."]
+            return []
     def chain_of_thought(self, doc=False, hint=""):
         """
         Generates a prompt using the chain-of-thought strategy.
-        If 'doc' is True, it follows a detailed documentation-oriented prompt strategy based on the round number.
-        If 'doc' is False, it provides general guidance for early round numbers and focuses on HTTP methods for later rounds.
 
         Args:
             doc (bool): Determines whether the documentation-oriented chain of thought should be used.
@@ -126,70 +158,30 @@ def chain_of_thought(self, doc=False, hint=""):
             "Make the OpenAPI specification available to developers by incorporating it into your API documentation site and keep the documentation up to date with API changes."
         ]
 
-        http_methods = [  "PUT", "DELETE"]
-        http_phase = {
-            5: http_methods[0],
-            10: http_methods[1]
-        }
-
+        http_methods = ["PUT", "DELETE"]
+        http_phase = {10: http_methods[0], 15: http_methods[1]}
         if doc:
-            if self.round < 5:
-
-                chain_of_thought_steps = [
-                                             f"Identify all available endpoints via GET Requests. Exclude those in this list: {self.found_endpoints}", f"Note down the response structures, status codes, and headers for each endpoint.",
-                                             f"For each endpoint, document the following details: URL, HTTP method, "
-                                             f"query parameters and path variables, expected request body structure for  requests, response structure for successful and error responses."
-                                         ] + common_steps
+            if self.round <= 5:
+                chain_of_thought_steps = self.get_initial_steps(common_steps)
+            elif self.round <= 10:
+                phase = http_phase.get(min(filter(lambda x: self.round <= x, http_phase.keys())))
+                chain_of_thought_steps = self.get_phase_steps(phase, common_steps)
             else:
-                if self.round <= 10:
-                    phase = http_phase.get(min(filter(lambda x: self.round <= x, http_phase.keys())))
-                    print(f'phase:{phase}')
-                    if phase != "DELETE":
-                        chain_of_thought_steps = [
-                                             f"Identify for all endpoints {self.found_endpoints} excluding {self.endpoint_found_methods[phase]} a valid HTTP method {phase} call.",
-                                             self.get_http_action_template(phase)
-                                         ] + common_steps
-                    else:
-                        chain_of_thought_steps = [
-                                                 f"Check for all endpoints the DELETE method. Delete the first instance for all endpoints. ",
-                                                 self.get_http_action_template(phase)
-                                             ] + common_steps
-                else:
-                    endpoints_needing_help = []
-                    endpoints_and_needed_methods = {}
-
-                    # Standard HTTP methods
-                    http_methods = {"GET", "POST", "PUT", "DELETE"}
-
-                    for endpoint in self.endpoint_methods:
-                        # Calculate the missing methods for the current endpoint
-                        missing_methods = http_methods - set(self.endpoint_methods[endpoint])
-
-                        if len(self.endpoint_methods[endpoint]) < 4:
-                            endpoints_needing_help.append(endpoint)
-                            # Add the missing methods to the dictionary
-                            endpoints_and_needed_methods[endpoint] = list(missing_methods)
-
-                    print(f'endpoints_and_needed_methods: {endpoints_and_needed_methods}')
-                    print(f'first endpoint in list: {endpoints_needing_help[0]}')
-                    print(f'methods needed for first endpoint: {endpoints_and_needed_methods[endpoints_needing_help[0]][0]}')
-
-                    chain_of_thought_steps = [f"For enpoint {endpoints_needing_help[0]} find this missing method :{endpoints_and_needed_methods[endpoints_needing_help[0]][0]} "
-                                              f"If all the HTTP methods have already been found for an endpoint, then do not include this endpoint in your search. ",]
-
+                chain_of_thought_steps = self.get_endpoints_needing_help()
         else:
             if self.round == 0:
-                chain_of_thought_steps = ["Let's think step by step."]  # Zero shot prompt
+                chain_of_thought_steps = ["Let's think step by step."]
             elif self.round <= 20:
                 focus_phases = ["endpoints", "HTTP method GET", "HTTP method POST and PUT", "HTTP method DELETE"]
                 focus_phase = focus_phases[self.round // 5]
                 chain_of_thought_steps = [f"Just focus on the {focus_phase} for now."]
             else:
                 chain_of_thought_steps = ["Look for exploits."]
 
-        print(f'chain of thought steps: {chain_of_thought_steps}')
-        prompt = self.check_prompt(self.previous_prompt,
-                                   chain_of_thought_steps + [hint] if hint else chain_of_thought_steps)
+        if hint:
+            chain_of_thought_steps.append(hint)
+
+        prompt = self.check_prompt(self.previous_prompt, chain_of_thought_steps)
         return prompt
 
     def token_count(self, text):

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -9,7 +9,7 @@
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
 from hackingBuddyGPT.capabilities.record_note import RecordNote
 from hackingBuddyGPT.usecases.agents import Agent
-from hackingBuddyGPT.usecases.web_api_testing.utils.documentation_handler import DocumentationHandler
+from hackingBuddyGPT.usecases.web_api_testing.utils.openapi_specification_manager import OpenAPISpecificationManager
 from hackingBuddyGPT.usecases.web_api_testing.utils.llm_handler import LLMHandler
 from hackingBuddyGPT.usecases.web_api_testing.prompt_engineer import PromptEngineer, PromptStrategy
 from hackingBuddyGPT.usecases.web_api_testing.utils.response_handler import ResponseHandler
@@ -52,7 +52,7 @@ def init(self):
         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)
+        self.documentation_handler = OpenAPISpecificationManager(self.llm_handler, self.response_handler)
 
     def _setup_capabilities(self):
         notes = self._context["notes"]
@@ -74,7 +74,7 @@ def _setup_initial_prompt(self):
                                               response_handler=self.response_handler)
 
 
-    def all_http_methods_found(self):
+    def all_http_methods_found(self,turn):
         print(f'found endpoints:{self.documentation_handler.endpoint_methods.items()}')
         print(f'found endpoints values:{self.documentation_handler.endpoint_methods.values()}')
 
@@ -83,17 +83,20 @@ def all_http_methods_found(self):
         print(f'found endpoints:{found_endpoints}')
         print(f'expected endpoints:{expected_endpoints}')
         print(f'correct? {found_endpoints== expected_endpoints}')
-        if found_endpoints== expected_endpoints or found_endpoints == expected_endpoints -1:
+        if found_endpoints > 0 and (found_endpoints== expected_endpoints) :
             return True
         else:
+            if turn == 20:
+                if found_endpoints > 0 and (found_endpoints == expected_endpoints):
+                    return True
             return False
 
     def perform_round(self, turn: int):
         prompt = self.prompt_engineer.generate_prompt(doc=True)
         response, completion = self.llm_handler.call_llm(prompt)
-        return self._handle_response(completion, response)
+        return self._handle_response(completion, response, turn)
 
-    def _handle_response(self, completion, response):
+    def _handle_response(self, completion, response, turn):
         message = completion.choices[0].message
         tool_call_id = message.tool_calls[0].id
         command = pydantic_core.to_json(response).decode()
@@ -106,7 +109,6 @@ def _handle_response(self, completion, response):
             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", "404" ,"Client Error: Not Found"]
-            print(f'result_str:{result_str}')
             if not result_str in invalid_flags  or any(item in result_str for item in invalid_flags):
                 self.prompt_engineer.found_endpoints = self.documentation_handler.update_openapi_spec(response, result)
                 self.documentation_handler.write_openapi_to_yaml()
@@ -120,8 +122,7 @@ def _handle_response(self, completion, response):
                         http_methods_dict[method].append(endpoint)
                 self.prompt_engineer.endpoint_found_methods =  http_methods_dict
                 self.prompt_engineer.endpoint_methods = self.documentation_handler.endpoint_methods
-                print(f'SCHEMAS:{self.prompt_engineer.schemas}')
-        return self.all_http_methods_found()
+        return self.all_http_methods_found(turn)
 
 
 

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

@@ -1,4 +1,4 @@
-from .documentation_handler import DocumentationHandler
+from .openapi_specification_manager import OpenAPISpecificationManager
 from .llm_handler import LLMHandler
 from .response_handler import ResponseHandler
 from .openapi_parser import OpenAPISpecificationParser

src/hackingBuddyGPT/usecases/web_api_testing/utils/documentation_handler.py renamed to src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_specification_manager.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from hackingBuddyGPT.capabilities.yamlFile import YAMLFile
 
-class DocumentationHandler:
+class OpenAPISpecificationManager:
     """
     Handles the generation and updating of an OpenAPI specification document based on dynamic API responses.
 
@@ -51,7 +51,7 @@ def __init__(self, llm_handler, response_handler):
             "yaml": YAMLFile()
         }
 
-    def partial_match(self, element, string_list):
+    def is_partial_match(self, element, string_list):
         return any(element in string or string in element for string in string_list)
 
     def update_openapi_spec(self, resp, result):
@@ -66,7 +66,7 @@ def update_openapi_spec(self, resp, result):
 
         if request.__class__.__name__ == 'RecordNote':  # TODO: check why isinstance does not work
             self.check_openapi_spec(resp)
-        if request.__class__.__name__ == 'HTTPRequest':
+        elif request.__class__.__name__ == 'HTTPRequest':
             path = request.path
             method = request.method
             print(f'method: {method}')
@@ -107,7 +107,7 @@ def update_openapi_spec(self, resp, result):
 
                     if '1' not in path and x != "":
                         endpoint_methods[path].append(method)
-                    elif self.partial_match(x, endpoints.keys()):
+                    elif self.is_partial_match(x, endpoints.keys()):
                         path = f"/{x}"
                         print(f'endpoint methods = {endpoint_methods}')
                         print(f'new path:{path}')

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

@@ -50,11 +50,14 @@ def parse_http_status_line(self, status_line):
         """
         if status_line == "Not a valid HTTP method":
             return status_line
-        if status_line and " " in status_line:
-            protocol, status_code, status_message = status_line.split(' ', 2)
-            status_message = status_message.split("\r\n")[0]
+        status_line = status_line.split('\r\n')[0]
+        # Regular expression to match valid HTTP status lines
+        match = re.match(r'^(HTTP/\d\.\d) (\d{3}) (.*)$', status_line)
+        if match:
+            protocol, status_code, status_message = match.groups()
             return f'{status_code} {status_message}'
-        raise ValueError("Invalid HTTP status line")
+        else:
+            raise ValueError("Invalid HTTP status line")
 
     def extract_response_example(self, html_content):
         """