Adjusted tests and refactored code for better readability

src/hackingBuddyGPT/usecases/web_api_testing/prompt_generation/information/pentesting_information.py

@@ -69,7 +69,7 @@ def __init__(self, openapi_spec_parser: OpenAPISpecificationParser, config) -> N
                                      PromptPurpose.VERIY_SETUP,
                                      #PromptPurpose.AUTHENTICATION,
                                      #PromptPurpose.AUTHORIZATION,
-                                     PromptPurpose.SPECIAL_AUTHENTICATION,
+                                     #PromptPurpose.SPECIAL_AUTHENTICATION,
                                      PromptPurpose.INPUT_VALIDATION,
                                      PromptPurpose.ERROR_HANDLING_INFORMATION_LEAKAGE,
                                      PromptPurpose.SESSION_MANAGEMENT,
@@ -280,8 +280,8 @@ def verify_setup(self) -> List:
                                 "objective": "Get Valid token",
                                 "steps": [
                                     f"Endpoint to use : {login_path}\n"
-                                    f"Send a POST request to the {login_schema} with the correct credentials of user:{account}.\n"
-                                    f"Request body should be in application/json and look similar to this: {{ {account}}}"
+                                    f"Send a POST request to the {login_path} with the correct credentials of user:{account}.\n"
+                                    f"Request body should be in application/json and look similar to this: {{ {login_schema}}}"
                                 ],
                                 "path": [login_path],
                                 "token": [account.get("token")],
@@ -1968,10 +1968,10 @@ def test_brute_force(self, login_path, login_schema, prompts, number=100):
                             "objective": "Brute Force Attack",
                             "steps": str,
                             "expected_response_code":
-                                [[
+                                [
                                     "429 Too Many Requests to indicate effective rate-limiting measures against brute force attacks.",
                                     "403 Forbidden as an alternative response indicating account lockout or other security measures preventing further attempts."
-                                ]] * 10,
+                                ] * 10,
                             "token": [""] * 10,
                             "path": [login_path] * 10,
                             "security":
@@ -1982,21 +1982,22 @@ def test_brute_force(self, login_path, login_schema, prompts, number=100):
                         }
                     )
 
-            change_email = [ep for ep in self.endpoints if "change-email"]
-            for endpoint in change_email:
+            change_email = self.extract_change_email_endpoints_with_schema()
+            for endpoint, schema in change_email:
                 for account in self.accounts:
                     if isinstance(endpoint, dict):
                         schema = endpoint.get("schema")
                         schema = self.fill_schema(schema)
-
+                    else:
+                        schema = self.fill_schema(schema)
                         prompts.append(  # Insecure Email Change Process
                             {
                                 "objective": "Insecure Email Change Process",
-                                "steps": f"Send a POST request to endpoint {endpoint} with schema {schema}",
+                                "steps": [f"Send a POST request to endpoint {endpoint} with schema {schema}"],
                                 "expected_response_code":
-                                    [[
+                                    [
                                         "403 Forbidden if proper authorization checks prevent changing another user's email."
-                                    ]],
+                                    ],
                                 "token": [account.get("token")],
                                 "path": [endpoint],
                                 "security":
@@ -2024,11 +2025,12 @@ def test_brute_force(self, login_path, login_schema, prompts, number=100):
 
             for reset in reset_password:
                 user = {"email": "adam007@example.com"}
+                # triggers otp being sent
+                new_otp_calls = f"Send a POST request to endpoint{reset} for victim user:{user}.\n" +  otp_calls
                 prompts.append(  # OTP Bruteforce Protection
                     {
                         "objective": "OTP Bruteforce Protection",
-                        "steps": f"Send a POST request to endpoint{reset} for victim user:{user}.\n" +  # triggers otp being sent
-                                 otp_calls,
+                        "steps": [new_otp_calls],
                         "expected_response_code":
                             [
                                 "429 Too Many Requests indicating effective rate-limiting against OTP bruteforce attempts.",
@@ -2046,6 +2048,24 @@ def test_brute_force(self, login_path, login_schema, prompts, number=100):
 
         return prompts
 
+    def extract_change_email_endpoints_with_schema(self) -> list[tuple]:
+        """
+        Extracts all endpoints containing 'change-email' and returns (path, schema) pairs.
+
+        Returns:
+            list: List of tuples (path, schema) for change-email endpoints.
+        """
+        result = []
+
+        for path, methods in self.endpoints.items():
+            if "change-email" in path:
+                for method, method_detail in methods.items():
+                    schema = method_detail.get("requestBody", {}).get("content", {}).get("application/json", {}).get(
+                        "schema")
+                    result.append((path, schema))
+
+        return result
+
     def test_css(self, endpoint, prompts, schema=None):
         if schema:
             prompts.append(

src/hackingBuddyGPT/usecases/web_api_testing/prompt_generation/prompts/task_planning/task_planning_prompt.py

@@ -125,12 +125,20 @@ def _get_pentesting_steps(self, move_type: str, common_step: Optional[str] = "")
                         self.explored_steps.append(task_planning_test_case)
 
 
+
                         self.prompt_helper.current_user = self.prompt_helper.get_user_from_prompt(self.current_sub_step,
                                                                                                   self.pentesting_information.accounts)
+
                         self.prompt_helper.counter = self.counter
 
                         step = self.transform_test_case_to_string(self.current_step, "steps")
+
+                        if self.prompt_helper.current_user is not None or isinstance(self.prompt_helper.current_user,                                                      dict):
+                            if "token" in self.prompt_helper.current_user and "'{{token}}'" in step:
+                                step = step.replace("'{{token}}'", self.prompt_helper.current_user.get("token"))
+
                         print(f'sub step:{self.current_sub_step}')
+                        print(f'step:{step}')
                         print(f'purpose:{self.purpose}')
                         self.counter += 1
                         # if last step of exploration, change purpose to next

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

@@ -49,8 +49,12 @@ def parse_http_response(self, raw_response: str) -> Tuple[Optional[int], Dict[st
 
         if body != {} and bool(body and not body.isspace()):
             body = json.loads(body)[0]
-        else:
-            body = "Empty"
+
+        if body == "":
+            for line in header_lines:
+                if line.startswith("{") or line.startswith("["):
+                    body = line
+                    body = json.loads(body)
 
         status_line = header_lines[0].strip()
         headers = {
@@ -91,7 +95,7 @@ def analyze_parsed_response(
             Optional[Dict[str, Any]]: The analysis results based on the purpose.
         """
         analysis_methods = {
-            PromptPurpose.AUTHENTICATION_AUTHORIZATION: self.analyze_authentication_authorization(
+            PromptPurpose.AUTHENTICATION: self.analyze_authentication_authorization(
                 status_code, headers, body
             ),
             PromptPurpose.INPUT_VALIDATION: self.analyze_input_validation(status_code, headers, body),

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

@@ -105,7 +105,7 @@ def parse_http_response(self, raw_response: str):
         body = header_body_split[1] if len(header_body_split) > 1 else ""
         if body == "":
             for line in header_lines:
-                if line.startswith("{"):
+                if line.startswith("{") or line.startswith("["):
                     body = line
 
         status_line = header_lines[0].strip()
@@ -135,8 +135,7 @@ def parse_http_response(self, raw_response: str):
                         self.prompt_helper.current_user["token"] = body["token"]
                         self.token = body["token"]
                         for account in self.prompt_helper.accounts:
-                                if account.get("x") == self.prompt_helper.current_user.get(
-                                        "x"):
+                                if account.get("x") == self.prompt_helper.current_user.get("x"):
                                     if  "token" not in account.keys():
                                         account["token"] = self.token
                                     else:

tests/test_response_analyzer.py

@@ -1,68 +1,112 @@
+import json
 import unittest
-from unittest.mock import patch
-
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import (
-    PromptPurpose,
-)
-from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_analyzer import (
-    ResponseAnalyzer,
-)
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptPurpose
+from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_analyzer import ResponseAnalyzer
 
 
 class TestResponseAnalyzer(unittest.TestCase):
+
     def setUp(self):
-        # Example HTTP response to use in tests
-        self.raw_http_response = """HTTP/1.1 404 Not Found
-        Date: Fri, 16 Aug 2024 10:01:19 GMT
-        Content-Type: application/json; charset=utf-8
-        Content-Length: 2
-        Connection: keep-alive
-        Report-To: {"group":"heroku-nel","max_age":3600,"endpoints":[{"url":"https://nel.heroku.com/reports?ts=1723802269&sid=e11707d5-02a7-43ef-b45e-2cf4d2036f7d&s=dkvm744qehjJmab8kgf%2BGuZA8g%2FCCIkfoYc1UdYuZMc%3D"}]}
-        X-Powered-By: Express
-        X-Ratelimit-Limit: 1000
-        X-Ratelimit-Remaining: 999
-        X-Ratelimit-Reset: 1723802321
-        Cache-Control: max-age=43200
-        Server: cloudflare
-
-        {}"""
-
-    def test_parse_http_response(self):
+        self.auth_headers = (
+            "HTTP/1.1 200 OK\n"
+            "Authorization: Bearer token\n"
+            "X-Ratelimit-Limit: 1000\n"
+            "X-Ratelimit-Remaining: 998\n"
+            "X-Ratelimit-Reset: 1723802321\n"
+            "\n"
+            '[{"message": "Welcome!"}]'
+        )
+
+        self.error_body = (
+            "HTTP/1.1 403 Forbidden\n"
+            "Content-Type: application/json\n"
+            "\n"
+            '[{"error": "Access denied"}]'
+        )
+
+        self.validation_fail = (
+            "HTTP/1.1 400 Bad Request\n"
+            "X-Content-Type-Options: nosniff\n"
+            "\n"
+            '[{"error": "Invalid input"}]'
+        )
+
+    def test_parse_http_response_success(self):
+        analyzer = ResponseAnalyzer()
+        status, headers, body = analyzer.parse_http_response(self.auth_headers)
+
+        self.assertEqual(200, status)
+        self.assertIn("Authorization", headers)
+        if isinstance(body, dict):
+            msg = body.get("message")
+            self.assertEqual( "Welcome!",msg )
+
+    def test_parse_http_response_invalid(self):
         analyzer = ResponseAnalyzer()
-        status_code, headers, body = analyzer.parse_http_response(self.raw_http_response)
+        status, headers, body = analyzer.parse_http_response(self.error_body)
 
-        self.assertEqual(status_code, 404)
-        self.assertEqual(headers["Content-Type"], "application/json; charset=utf-8")
-        self.assertEqual(body, "Empty")
+        self.assertEqual( 403, status)
+        if isinstance(body, dict):
+            msg = body.get("message")
 
-    def test_analyze_authentication_authorization(self):
-        analyzer = ResponseAnalyzer(PromptPurpose.AUTHENTICATION_AUTHORIZATION)
-        analysis = analyzer.analyze_response(self.raw_http_response)
+            self.assertEqual( "Access denied", msg)
 
-        self.assertEqual(analysis["status_code"], 404)
-        self.assertEqual(analysis["authentication_status"], "Unknown")
-        self.assertTrue(analysis["content_body"], "Empty")
-        self.assertIn("X-Ratelimit-Limit", analysis["rate_limiting"])
+    def test_analyze_authentication(self):
+        analyzer = ResponseAnalyzer(PromptPurpose.AUTHENTICATION)
+        result = analyzer.analyze_response(self.auth_headers)
 
-    def test_analyze_input_validation(self):
+        self.assertEqual(result["status_code"], 200)
+        self.assertEqual(result["authentication_status"], "Authenticated")
+        self.assertTrue(result["auth_headers_present"])
+        self.assertIn("X-Ratelimit-Limit", result["rate_limiting"])
+
+    def test_analyze_input_validation_invalid(self):
         analyzer = ResponseAnalyzer(PromptPurpose.INPUT_VALIDATION)
-        analysis = analyzer.analyze_response(self.raw_http_response)
+        result = analyzer.analyze_response(self.validation_fail)
+
+        self.assertEqual(result["status_code"], 400)
+        self.assertEqual(result["is_valid_response"], "Invalid")
+        self.assertTrue(result["security_headers_present"])
 
-        self.assertEqual(analysis["status_code"], 404)
-        self.assertEqual(analysis["is_valid_response"], "Error")
-        self.assertTrue(analysis["response_body"], "Empty")
-        self.assertIn("security_headers_present", analysis)
+    def test_is_valid_input_response(self):
+        analyzer = ResponseAnalyzer()
+        self.assertEqual(analyzer.is_valid_input_response(200, "data"), "Valid")
+        self.assertEqual(analyzer.is_valid_input_response(400, "error"), "Invalid")
+        self.assertEqual(analyzer.is_valid_input_response(500, "error"), "Error")
+        self.assertEqual(analyzer.is_valid_input_response(999, "???"), "Unexpected")
 
-    @patch("builtins.print")
-    def test_print_analysis(self, mock_print):
+    def test_document_findings(self):
+        analyzer = ResponseAnalyzer()
+        document = analyzer.document_findings(
+            status_code=403,
+            headers={"Content-Type": "application/json"},
+            body="Access denied",
+            expected_behavior="Access should be allowed",
+            actual_behavior="Access denied"
+        )
+        self.assertEqual(document["Status Code"], 403)
+        self.assertIn("Access denied", document["Actual Behavior"])
+
+    def test_print_analysis_output_structure(self):
         analyzer = ResponseAnalyzer(PromptPurpose.INPUT_VALIDATION)
-        analysis = analyzer.analyze_response(self.raw_http_response)
-        analysis_str = analyzer.print_analysis(analysis)
+        result = analyzer.analyze_response(self.validation_fail)
+        printed = analyzer.print_analysis(result)
+
+        self.assertIn("HTTP Status Code: 400", printed)
+        self.assertIn("Valid Response: Invalid", printed)
+        self.assertIn("Security Headers Present", printed)
 
-        # Check that the correct calls were made to print
-        self.assertIn("HTTP Status Code: 404", analysis_str)
-        self.assertIn("Response Body: Empty", analysis_str)
-        self.assertIn("Security Headers Present: Yes", analysis_str)
+    def test_report_issues_found(self):
+        analyzer = ResponseAnalyzer()
+        document = analyzer.document_findings(
+            status_code=200,
+            headers={},
+            body="test",
+            expected_behavior="User not authenticated",
+            actual_behavior="User is authenticated"
+        )
+        # Just ensure no exceptions, prints okay
+        analyzer.report_issues(document)
 
 
 if __name__ == "__main__":