Merge remote-tracking branch 'refs/remotes/origin/web_api_testing' into development_without_spacy

Author: DianaStrauss

src/hackingBuddyGPT/capabilities/submit_http_method.py

@@ -16,6 +16,7 @@ class SubmitHTTPMethod(Capability):
     _client = requests.Session()
     host: str
     follow_redirects: bool = False
+    success_function: Callable[[], None] = None
 
 
     submitted_valid_http_methods: Set[str] = field(default_factory=set, init=False)
@@ -67,7 +68,11 @@ def __call__(self, method: Literal["GET", "HEAD", "POST", "PUT", "DELETE", "OPTI
             return str(e)
 
         headers = "\r\n".join(f"{k}: {v}" for k, v in resp.headers.items())
-
+        if len(self.submitted_valid_http_methods) == len(self.valid_http_methods):
+            if self.success_function is not None:
+                self.success_function()
+            else:
+                return "All methods submitted, congratulations"
         # turn the response into "plain text format" for responding to the prompt
         return f"HTTP/1.1 {resp.status_code} {resp.reason}\r\n{headers}\r\n\r\n{resp.text}"""
 

src/hackingBuddyGPT/cli/wintermute.py

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

src/hackingBuddyGPT/usecases/web_api_testing/prompt_engineer.py

@@ -32,7 +32,9 @@ def __init__(self, strategy, llm_handler, history, schemas, response_handler):
         self.response_handler = response_handler
         self.llm_handler = llm_handler
         self.round = 0
-        self.found_endpoints = []
+        self.found_endpoints = ["/"]
+        self.endpoint_methods = {}
+        self.endpoint_found_methods = {}
         model_name = "en_core_web_sm"
 
         # Check if the model is already installed
@@ -105,7 +107,7 @@ def get_http_action_template(self, method):
 
         else:
             return (
-                f"Create HTTPRequests of type {method} considering the found schemas: {self.schemas} and understand the responses. Ensure that they are correct requests.")
+                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 chain_of_thought(self, doc=False, hint=""):
@@ -129,34 +131,56 @@ 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 = [ "POST", "DELETE", "PUT"]
+        http_methods = [  "PUT", "DELETE"]
         http_phase = {
-            6: http_methods[0],
-            16: http_methods[1], # Delete one of instance of this:{self.llm_handler.get_created_objects()}",
-            20: http_methods[2]
+            5: http_methods[0],
+            10: http_methods[1]
         }
 
         if doc:
-            if self.round <= 5:
+            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
             else:
-                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 all valid calls for HTTP method {phase}.",
+                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 = [
+                    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. ",]
 
         else:
             if self.round == 0:
@@ -187,6 +211,7 @@ def token_count(self, text):
         doc = self.nlp(text)
         # Count tokens that aren't punctuation marks
         tokens = [token for token in doc if not token.is_punct]
+        print(f'TOKENS: {len(tokens)}')
         return len(tokens)
 
 

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -75,13 +75,23 @@ def _setup_initial_prompt(self):
 
 
     def all_http_methods_found(self):
-        self._log.console.print(Panel("All HTTP methods found! Congratulations!", title="system"))
-        self._all_http_methods_found = True
+        print(f'found endpoints:{self.documentation_handler.endpoint_methods.items()}')
+        print(f'found endpoints values:{self.documentation_handler.endpoint_methods.values()}')
+
+        found_endpoints = sum(len(value_list) for value_list in self.documentation_handler.endpoint_methods.values())
+        expected_endpoints = len(self.documentation_handler.endpoint_methods.keys())*4
+        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:
+            return True
+        else:
+            return False
 
     def perform_round(self, turn: int):
         prompt = self.prompt_engineer.generate_prompt(doc=True)
         response, completion = self.llm_handler.call_llm(prompt)
-        self._handle_response(completion, response)
+        return self._handle_response(completion, response)
 
     def _handle_response(self, completion, response):
         message = completion.choices[0].message
@@ -101,8 +111,17 @@ def _handle_response(self, completion, response):
                 self.prompt_engineer.found_endpoints = self.documentation_handler.update_openapi_spec(response, result)
                 self.documentation_handler.write_openapi_to_yaml()
                 self.prompt_engineer.schemas = self.documentation_handler.schemas
+                from collections import defaultdict
+                http_methods_dict = defaultdict(list)
+
+                # Iterate through the original dictionary
+                for endpoint, methods in self.documentation_handler.endpoint_methods.items():
+                    for method in methods:
+                        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()
 
 
 

src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py

@@ -95,7 +95,7 @@ def _setup_capabilities(self):
         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, methods_set, self.host),
+            "submit_http_method": HTTPRequest(self.host),
             "http_request": HTTPRequest(self.host),
             "record_note": RecordNote(notes)
         }
@@ -134,8 +134,7 @@ 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))
 
-        return self._all_http_methods_found
-
+        return self.all_http_methods_found()
 @use_case("Minimal implementation of a web API testing use case")
 class SimpleWebAPITestingUseCase(AutonomousAgentUseCase[SimpleWebAPITesting]):
     pass

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

@@ -29,6 +29,7 @@ def __init__(self, llm_handler, response_handler):
         """
         self.response_handler = response_handler
         self.schemas = {}
+        self.endpoint_methods ={}
         self.filename = f"openapi_spec_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.yaml"
         self.openapi_spec = {
             "openapi": "3.0.0",
@@ -50,6 +51,9 @@ def __init__(self, llm_handler, response_handler):
             "yaml": YAMLFile()
         }
 
+    def 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):
         """
         Updates the OpenAPI specification based on the API response provided.
@@ -67,31 +71,51 @@ def update_openapi_spec(self, resp, result):
             method = request.method
             print(f'method: {method}')
             # Ensure that path and method are not None and method has no numeric characters
+            # Ensure path and method are valid and method has no numeric characters
             if path and method:
+                endpoint_methods = self.endpoint_methods
+                endpoints = self.openapi_spec['endpoints']
+                x = path.split('/')[1]
+
                 # Initialize the path if not already present
-                if path not in self.openapi_spec['endpoints']:
-                    self.openapi_spec['endpoints'][path] = {}
+                if path not in endpoints and x != "":
+                    endpoints[path] = {}
+                    if '1' not in path:
+                        endpoint_methods[path] = []
+
                 # Update the method description within the path
-                example, reference, self.openapi_spec = self.response_handler.parse_http_response_to_openapi_example(self.openapi_spec, result, path, method)
+                example, reference, self.openapi_spec = self.response_handler.parse_http_response_to_openapi_example(
+                    self.openapi_spec, result, path, method
+                )
                 self.schemas = self.openapi_spec["components"]["schemas"]
-                if example is not None or reference is not None:
-                    self.openapi_spec['endpoints'][path][method.lower()] = {
+
+                if example or reference:
+                    endpoints[path][method.lower()] = {
                         "summary": f"{method} operation on {path}",
                         "responses": {
                             "200": {
                                 "description": "Successful response",
                                 "content": {
                                     "application/json": {
-                                        "schema": {
-                                            "$ref": reference
-                                        },
+                                        "schema": {"$ref": reference},
                                         "examples": example
                                     }
                                 }
                             }
                         }
                     }
-            return  list(self.openapi_spec['endpoints'].keys())
+
+                    if '1' not in path and x != "":
+                        endpoint_methods[path].append(method)
+                    elif self.partial_match(x, endpoints.keys()):
+                        path = f"/{x}"
+                        print(f'endpoint methods = {endpoint_methods}')
+                        print(f'new path:{path}')
+                        endpoint_methods[path].append(method)
+
+                    endpoint_methods[path] = list(set(endpoint_methods[path]))
+
+            return list(endpoints.keys())
 
     def write_openapi_to_yaml(self):
         """

tests/test_web_api_documentation.py

@@ -0,0 +1,43 @@
+import argparse
+import unittest
+from hackingBuddyGPT.usecases.base import use_cases
+from hackingBuddyGPT.usecases.web_api_testing.simple_web_api_testing import SimpleWebAPITestingUseCase
+from hackingBuddyGPT.utils import DbStorage, Console
+
+
+
+class WebAPIDocumentationTestCase(unittest.TestCase):
+    def test_simple_web_api_testing(self):
+
+        log_db = DbStorage(':memory:')
+        console = Console()
+
+        log_db.init()
+        parser = argparse.ArgumentParser()
+        subparser = parser.add_subparsers(required=True)
+        for name, use_case in use_cases.items():
+            use_case.build_parser(subparser.add_parser(
+                name=use_case.name,
+                help=use_case.description
+            ))
+
+        parsed = parser.parse_args(["SimpleWebAPIDocumentation"])
+        instance = parsed.use_case(parsed)
+
+        agent = instance.agent
+        simple_web_api_documentation = SimpleWebAPITestingUseCase(
+            agent=agent,
+            log_db=log_db,
+            console=console,
+            tag='web_api_documentation',
+            max_turns=20
+        )
+
+        simple_web_api_documentation.init()
+        result = simple_web_api_documentation.run()
+        print(f'result: {result}')
+        assert result is True
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_web_api_testing.py

@@ -0,0 +1,42 @@
+import argparse
+import unittest
+from hackingBuddyGPT.usecases.base import use_cases
+from hackingBuddyGPT.usecases.web_api_testing.simple_web_api_testing import SimpleWebAPITestingUseCase
+from hackingBuddyGPT.utils import DbStorage, Console
+
+
+
+class WebAPITestingTestCase(unittest.TestCase):
+    def test_simple_web_api_testing(self):
+
+        log_db = DbStorage(':memory:')
+        console = Console()
+
+        log_db.init()
+        parser = argparse.ArgumentParser()
+        subparser = parser.add_subparsers(required=True)
+        for name, use_case in use_cases.items():
+            use_case.build_parser(subparser.add_parser(
+                name=use_case.name,
+                help=use_case.description
+            ))
+
+        parsed = parser.parse_args(["SimpleWebAPITesting"])
+        instance = parsed.use_case(parsed)
+
+        agent = instance.agent
+        simple_web_api_testing = SimpleWebAPITestingUseCase(
+            agent=agent,
+            log_db=log_db,
+            console=console,
+            tag='web_api_testing',
+            max_turns=20
+        )
+
+        simple_web_api_testing.init()
+        result = simple_web_api_testing.run()
+        # TODO: find condition for testing
+
+
+if __name__ == '__main__':
+    unittest.main()