finished mocking web_api_documentation testing

Author: Diana Strauss

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -29,19 +29,19 @@ class SimpleWebAPIDocumentation(Agent):
     _all_http_methods_found: bool = False
 
     # Description for expected HTTP methods
-    http_method_description: str = parameter(
+    _http_method_description: str = parameter(
         desc="Pattern description for expected HTTP methods in the API response",
         default="A string that represents an HTTP method (e.g., 'GET', 'POST', etc.)."
     )
 
     # Template for HTTP methods in API requests
-    http_method_template: str = parameter(
+    _http_method_template: str = parameter(
         desc="Template to format HTTP methods in API requests, with {method} replaced by actual HTTP method names.",
         default="{method}"
     )
 
     # List of expected HTTP methods
-    http_methods: str = parameter(
+    _http_methods: str = parameter(
         desc="Expected HTTP methods in the API, as a comma-separated list.",
         default="GET,POST,PUT,PATCH,DELETE"
     )

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

@@ -43,7 +43,7 @@ def __init__(self, llm_handler, response_handler):
             "components": {"schemas": {}}
         }
         self.llm_handler = llm_handler
-        self.api_key = llm_handler.llm.api_key
+        #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)

tests/test_web_api_documentation.py

@@ -1,43 +1,77 @@
-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
+from unittest.mock import MagicMock, patch
 
+from hackingBuddyGPT.usecases import SimpleWebAPITesting
+from hackingBuddyGPT.usecases.web import MinimalWebTesting
+from hackingBuddyGPT.usecases.web_api_testing.simple_openapi_documentation import SimpleWebAPIDocumentationUseCase
+from hackingBuddyGPT.utils import DbStorage, Console
+from hackingBuddyGPT.utils.openai.openai_lib import OpenAILib
 
 
-class WebAPIDocumentationTestCase(unittest.TestCase):
-    def test_simple_web_api_testing(self):
+class TestSimpleWebAPITesting(unittest.TestCase):
 
+    @patch('hackingBuddyGPT.utils.openai.openai_lib.OpenAILib')
+    def setUp(self, MockOpenAILib):
+        # Mock the OpenAILib instance
+        self.mock_llm = MockOpenAILib.return_value
         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,
+        self.agent = SimpleWebAPITesting(llm=self.mock_llm)
+        self.agent.init()
+        self.simple_api_testing = SimpleWebAPIDocumentationUseCase(
+            agent=self.agent,
             log_db=log_db,
             console=console,
-            tag='web_api_documentation',
-            max_turns=20
+            tag='integration_test_linuxprivesc',
+            max_turns=len(self.mock_llm.responses)
         )
+        self.simple_api_testing.init()
+
+    def test_initial_prompt(self):
+        # Test if the initial prompt is set correctly
+        self.assertIn("You're a seasoned API tester tasked with evaluating the REST APIs of a website hosted at https://jsonplaceholder.typicode.com. Your primary objective is to thoroughly explore the APIs, understanding their endpoints, parameters, and responses, and then devise and execute potential exploits. Be diligent in documenting your findings as you navigate through the APIs. Avoid resorting to brute-force tactics. All necessary information can be obtained from the API endpoints themselves. Remember, if you encounter an HTTP method (A string that represents an HTTP method (e.g., 'GET', 'POST', etc.).), promptly submit it as it is of utmost importance.", self.agent._prompt_history[0]['content'])
+
+    def test_all_flags_found(self):
+        # Mock console.print to suppress output during testing
+        with patch('rich.console.Console.print'):
+            self.agent.all_http_methods_found()
+            self.assertFalse(self.agent.all_http_methods_found())
+
+    @patch('time.perf_counter', side_effect=[1, 2])  # Mocking perf_counter for consistent timing
+    def test_perform_round(self, mock_perf_counter):
+        # Prepare mock responses
+        mock_response = MagicMock()
+        mock_completion = MagicMock()
+
+        # Setup completion response with mocked data
+        mock_completion.choices[0].message.content = "Mocked LLM response"
+        mock_completion.choices[0].message.tool_calls = [MagicMock(id="tool_call_1")]
+        mock_completion.usage.prompt_tokens = 10
+        mock_completion.usage.completion_tokens = 20
+
+        # Mock the OpenAI LLM response
+        self.agent.llm.instructor.chat.completions.create_with_completion.return_value = (
+        mock_response, mock_completion)
+
+        # Mock the tool execution result
+        mock_response.execute.return_value = "Mocked tool execution result"
+
+        # Perform the round
+        result = self.agent.perform_round(1)
+
+        # Assertions
+        self.assertFalse(result)  # No flags found in this round
+
+        # Check if the LLM was called with the correct parameters
+        mock_create_with_completion = self.agent.llm.instructor.chat.completions.create_with_completion
 
-        simple_web_api_documentation.init()
-        result = simple_web_api_documentation.run()
-        print(f'result: {result}')
-        assert result is True
+        # if it can be called multiple times, use assert_called
+        self.assertEqual( 2, mock_create_with_completion.call_count)
 
+        # Check if the prompt history was updated correctly
+        self.assertEqual(5, len(self.agent._prompt_history))  # Initial message + LLM response + tool message
 
 if __name__ == '__main__':
     unittest.main()