Updated Tree of thought so that documentation works like chain of tho…

…ught

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

@@ -74,7 +74,7 @@ def update_openapi_spec(self, resp, result, result_str):
         status_code, status_message = result_str.split(" ", 1)
 
         if request.__class__.__name__ == "RecordNote":  # TODO: check why isinstance does not work
-            self.check_openapi_spec(resp)
+            #self.check_openapi_spec(resp)
             return list(self.openapi_spec["endpoints"].keys())
 
         if request.__class__.__name__ == "HTTPRequest":
@@ -90,7 +90,7 @@ def update_openapi_spec(self, resp, result, result_str):
 
             # Extract the main part of the path for checking partial matches
             path_parts = path.split("/")
-            main_path = path_parts[1] if len(path_parts) > 1 else ""
+            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:
@@ -105,7 +105,9 @@ def update_openapi_spec(self, resp, result, result_str):
 
             # Add example and reference to the method's responses if available
             if example or reference:
-                endpoints[path][method.lower()] = {
+                if path in endpoints.keys() and  method.lower() not in endpoints[path].values():
+
+                    endpoints[path][method.lower()] = {
                     "summary": f"{method} operation on {path}",
                     "responses": {
                         f"{status_code}": {
@@ -114,17 +116,17 @@ def update_openapi_spec(self, resp, result, result_str):
                                 "application/json": {
                                     "schema": {"$ref": reference},
                                     "examples": example
+                                    }
                                 }
                             }
                         }
                     }
-                }
 
-                # Update endpoint methods for the path
-                endpoint_methods[path].append(method)
+                    # Update endpoint methods for the path
+                    endpoint_methods[path].append(method)
 
-                # Ensure uniqueness of methods for each path
-                endpoint_methods[path] = list(set(endpoint_methods[path]))
+                    # Ensure uniqueness of methods for each path
+                    endpoint_methods[path] = list(set(endpoint_methods[path]))
 
         return list(self.openapi_spec["endpoints"].keys())
 

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

@@ -40,87 +40,83 @@ def __init__(self, response_handler: ResponseHandler = None, schemas: dict = Non
 
     import re
 
-    def find_missing_endpoints(self, endpoints: list) -> list:
+    def find_missing_endpoint(self, endpoints: dict) -> str:
         """
-        Identifies and returns the actual missing endpoint paths.
+        Identifies and returns the first missing endpoint path found.
 
         Args:
-            endpoints (dict): A dictionary of endpoint paths (e.g., {'/resources/': {...}, '/resources/:id': {...}}).
+            endpoints (dict): A dictionary of endpoint paths (e.g., {'/resources': {...}, '/resources/:id': {...}}).
 
         Returns:
-            list: A list of missing endpoint paths.
-                  Example: ['/resources/:id', '/products/']
+            str: The first missing endpoint path found.
+                 Example: '/resources/:id' or '/products'
         """
         general_endpoints = set()
         parameterized_endpoints = set()
 
-        # Extract resource names and categorize them
+        # Extract resource names and categorize them using regex
         for endpoint in endpoints:
-            match = re.match(r'^/([^/]+)/?$', endpoint)  # Match general endpoints like /resources/
+            # Match both general and parameterized patterns and categorize them
+            match = re.match(r'^/([^/]+)(/|/:id)?$', endpoint)
             if match:
-                general_endpoints.add(match.group(1))
-
-            match = re.match(r'^/([^/]+)/:id$', endpoint)  # Match parameterized endpoints like /resources/:id
-            if match:
-                parameterized_endpoints.add(match.group(1))
-
-        # Find resources that are missing either general or parameterized endpoints
-        missing_endpoints = []
-        all_resources = general_endpoints | parameterized_endpoints
-
-        for resource in all_resources:
+                resource = match.group(1)
+                if match.group(2) == '/' or match.group(2) is None:
+                    general_endpoints.add(resource)
+                elif match.group(2) == '/:id':
+                    parameterized_endpoints.add(resource)
+
+        # Find missing endpoints during the comparison
+        for resource in parameterized_endpoints:
             if resource not in general_endpoints:
-                missing_endpoints.append(f'/{resource}/')
+                return f'/{resource}'
+        for resource in general_endpoints:
             if resource not in parameterized_endpoints:
-                missing_endpoints.append(f'/{resource}/:id')
+                return f'/{resource}/:id'
 
-            # If only one missing endpoint is needed, break early
-            if len(missing_endpoints) == 1:
-                break
-
-        return missing_endpoints
+        # Return an empty string if no missing endpoints are found
+        return ""
 
     def get_endpoints_needing_help(self, info=""):
         """
         Identifies endpoints that need additional HTTP methods and returns guidance for the first missing method.
 
+        Args:
+            info (str): Additional information to include in the response.
+
         Returns:
             list: A list containing guidance for the first missing method of the first endpoint that needs help.
         """
-        endpoints_needing_help = []
-        endpoints_and_needed_methods = {}
         http_methods_set = {"GET", "POST", "PUT", "DELETE"}
-
         for endpoint, methods in self.endpoint_methods.items():
-            if len(methods) >= 4:
-                continue
-
-            # the endpoint needs help
             missing_methods = http_methods_set - set(methods)
-            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]
-            print(F'{first_endpoint}: {needed_method}')
-            if ":id" in first_endpoint:
-                first_endpoint = first_endpoint.replace(":id", "1")
-            return [
-                info + "/n",
-                f"For endpoint {first_endpoint}, find this missing method: {needed_method}. "
-                #f"If all HTTP methods have already been found for an endpoint, do not include this endpoint in your search."
-            ]
-        else:
-            missing_endpoints = self.find_missing_endpoints(endpoints=self.found_endpoints)
-            if missing_endpoints:
-                needed_method = "GET"
+            if missing_methods:
+                needed_method = next(iter(missing_methods))
+                formatted_endpoint = endpoint.replace(":id", "1") if ":id" in endpoint else endpoint
                 return [
-                info + "/n",
-                f"For endpoint {missing_endpoints[0]}, find this missing method: {needed_method}. "
-                # f"If all HTTP methods have already been found for an endpoint, do not include this endpoint in your search."
+                    f"{info}\n",
+                    f"For endpoint {formatted_endpoint}, find this missing method: {needed_method}. "
                 ]
 
+        # If no endpoints need help, find missing endpoints and suggest "GET"
+        missing_endpoint = self.find_missing_endpoint(endpoints=self.found_endpoints)
+        print(f"------------------------------------")
+        print(f"------------------------------------")
+        print(f"------------------------------------")
+        print(f"{info}\n{missing_endpoint}")
+        print(f"------------------------------------")
+        print(f"------------------------------------")
+        print(f"------------------------------------")
+
+        if missing_endpoint != "":
+            formatted_endpoint = missing_endpoint.replace(":id", "1") if ":id" in missing_endpoint else \
+            missing_endpoint
+            return [
+                f"{info}\n",
+                f"For endpoint {formatted_endpoint}, find this missing method: GET. "
+            ]
+
+        return []
+
     def get_http_action_template(self, method):
         """
         Constructs a consistent HTTP action description based on the provided method.

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

@@ -128,6 +128,8 @@ def parse_http_response_to_openapi_example(
                     key = entry.get("title") or entry.get("name") or entry.get("id")
                     entry_dict[key] = {"value": entry}
                     self.llm_handler.add_created_object(entry_dict[key], object_name)
+                    if len(entry_dict) > 3:
+                        break
             else:
                 key = body_dict.get("title") or body_dict.get("name") or body_dict.get("id")
                 entry_dict[key] = {"value": body_dict}
@@ -161,6 +163,8 @@ def parse_http_response_to_schema(
         Returns:
             Tuple[str, str, Dict[str, Any]]: A tuple containing the reference, object name, and updated OpenAPI specification.
         """
+        if "/" not in path:
+            return None, None, openapi_spec
         object_name = path.split("/")[1].capitalize().rstrip("s")
         properties_dict = {}
 

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -67,7 +67,7 @@ def init(self):
         self._setup_capabilities()
         self.llm_handler = LLMHandler(self.llm, self._capabilities)
         self.response_handler = ResponseHandler(self.llm_handler)
-        self.strategy = PromptStrategy.IN_CONTEXT
+        self.strategy = PromptStrategy.TREE_OF_THOUGHT
         self.documentation_handler = OpenAPISpecificationHandler(self.llm_handler, self.response_handler, self.strategy)
         self._setup_initial_prompt()
 
@@ -132,31 +132,37 @@ def perform_round(self, turn: int) -> bool:
         if turn == 1:
             last_endpoint_found_x_steps_ago = 0
             new_endpoint_count = len(self.documentation_handler.endpoint_methods)
-            last_number_of_found_endpoints = 0
-            while (last_endpoint_found_x_steps_ago <= new_endpoint_count + 2
-                   and last_endpoint_found_x_steps_ago <= 5
-                   and not self.found_all_http_methods):
+            last_number_of_found_endpoints = len(self.prompt_engineer.prompt_helper.found_endpoints)
+
+            # Explore mode: search for new endpoints until conditions are met
+            while (
+                    last_endpoint_found_x_steps_ago <= new_endpoint_count + 2
+                    and last_endpoint_found_x_steps_ago <= 5
+                    and not self.found_all_http_methods
+            ):
                 self.run_documentation(turn, "explore")
 
-                # Check if new endpoints have been found
+                # Update endpoint counts
                 current_endpoint_count = len(self.prompt_engineer.prompt_helper.found_endpoints)
-                if last_number_of_found_endpoints == len(self.prompt_engineer.prompt_helper.found_endpoints):
+
+                if current_endpoint_count == last_number_of_found_endpoints:
                     last_endpoint_found_x_steps_ago += 1
                 else:
-                    last_endpoint_found_x_steps_ago = 0  # Reset if a new endpoint is found
+                    last_endpoint_found_x_steps_ago = 0
+                    last_number_of_found_endpoints = current_endpoint_count
 
-                # Update if new endpoint methods are discovered
-                if len(self.documentation_handler.endpoint_methods) > new_endpoint_count:
-                    new_endpoint_count = len(self.documentation_handler.endpoint_methods)
+                # Check if new methods have been discovered
+                updated_endpoint_count = len(self.documentation_handler.endpoint_methods)
+                if updated_endpoint_count > new_endpoint_count:
+                    new_endpoint_count = updated_endpoint_count
                     self.prompt_engineer.open_api_spec = self.documentation_handler.openapi_spec
 
-                last_number_of_found_endpoints = current_endpoint_count
-
         elif turn == 20:
-            # Continue until all endpoints needing help are addressed
+            # Exploit mode: refine endpoints until no further help is needed
             while self.prompt_engineer.prompt_helper.get_endpoints_needing_help():
                 self.run_documentation(turn, "exploit")
                 self.prompt_engineer.open_api_spec = self.documentation_handler.openapi_spec
+
         else:
             # For other turns, run documentation in exploit mode
             self.run_documentation(turn, "exploit")