fixed shortening of prompt

Author: DianaStrauss

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

@@ -52,11 +52,26 @@ def write_analysis_to_report(self, analysis: List[str], purpose: Enum) -> None:
             analysis (List[str]): The analysis data to be recorded.
             purpose (Enum): An enumeration that describes the purpose of the analysis.
         """
+        # Open the file in read mode to check if the purpose already exists
+        try:
+            with open(self.report_name, 'r') as report:
+                content = report.read()
+        except FileNotFoundError:
+            # If file does not exist, treat as if the purpose doesn't exist
+            content = ""
+
+        # Check if the purpose.name is already in the content
+        if purpose.name not in content:
+            with open(self.report_name, 'a') as report:
+                report.write(
+                    '-------------------------------------------------------------------------------------------\n')
+                report.write(f'{purpose.name}:\n')
+
+        # Write the analysis data
         with open(self.report_name, 'a') as report:
-            report.write(f'{purpose.name}:\n')
             for item in analysis:
-                for line in item.split("\n"):
-                    if "note recorded" in line:
-                        continue
-                    else:
-                        report.write(line + "\n")
+                lines = item.split("\n")
+                filtered_lines = [line for line in lines if "note recorded" not in line]
+                report.write("\n".join(filtered_lines) + "\n")
+
+

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

@@ -14,14 +14,15 @@ class PenTestingInformation:
         exploit_steps (dict): A dictionary mapping each PromptPurpose to a list of exploitation steps.
     """
 
-    def __init__(self, schemas: Dict[str, List[str]] = None) -> object:
+    def __init__(self, schemas: Dict[str, List[str]] = None, endpoints: Dict[str, List[str]] = None) -> object:
         """
         Initializes the PenTestingInformation with optional schemas.
 
         Args:
             schemas (dict, optional): A dictionary containing valid schemas for input validation tests. Default is None.
         """
         self.schemas = schemas if schemas is not None else {}
+        self.endpoints = endpoints if endpoints is not None else {}
         self.explore_steps = self.init_steps
 
     @property
@@ -40,19 +41,22 @@ def init_steps(self) -> Dict[PromptPurpose, List[str]]:
             ],
             PromptPurpose.INPUT_VALIDATION: [
                 f"Examine how the API handles various types of user input by crafting and sending a POST request with valid data. "
-                f"Valid schemas are {self.schemas}.",
+                f"Valid schemas are {self.schemas} and endpoints are {self.endpoints}.",
                 f"Examine how the API handles various types of user input by crafting and sending a POST request with invalid data. "
-                f"Valid schemas are {self.schemas}.",
+                f"Valid schemas are {self.schemas} and endpoints are {self.endpoints}.",
                 f"Examine how the API handles various types of user input by crafting and sending a POST request with edge data. "
-                f"Valid schemas are {self.schemas}.",
+                f"Valid schemas are {self.schemas} and endpoints are {self.endpoints}.",
                 f"Omit required fields to see if the API correctly identifies them as missing by sending a POST request. "
-                f"Valid schemas are {self.schemas}.",
+                f"Valid schemas are {self.schemas} and endpoints are {self.endpoints}.",
                 f"Include special characters or injection attempts (e.g., SQL Injection, XSS) to see if the API sanitizes inputs properly, "
-                f"by sending a POST request. Valid schemas are {self.schemas}.",
-                f"Send data that doesn’t match the expected format by a POST request. Valid data formats are {self.schemas}.",
+                f"by sending a POST request. "
+                f"Valid schemas are {self.schemas} and endpoints are {self.endpoints}.",
+                f"Send data that doesn’t match the expected format by a POST request. "
+                f"Valid data formats are {self.schemas} and endpoints are {self.endpoints}.",
                 "Check for proper error handling, response codes, and sanitization.",
                 "Attempt to exploit common vulnerabilities by injecting malicious inputs, such as SQL injection, NoSQL injection, "
-                "cross-site scripting, and other injection attacks. Evaluate whether the API properly validates, escapes, and sanitizes "
+                "cross-site scripting, and other injection attacks. ",
+                "Evaluate whether the API properly validates, escapes, and sanitizes "
                 "all user-supplied data, ensuring no unexpected behavior or security vulnerabilities are exposed."
             ],
             PromptPurpose.ERROR_HANDLING_INFORMATION_LEAKAGE: [

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

@@ -1,5 +1,6 @@
 from instructor.retry import InstructorRetryException
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy, PromptContext
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy, \
+    PromptContext, PromptPurpose
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompt_generation_helper import PromptGenerationHelper
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.task_planning import ChainOfThoughtPrompt, TreeOfThoughtPrompt
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.state_learning import InContextLearningPrompt
@@ -12,7 +13,7 @@ class PromptEngineer:
 
     def __init__(self, strategy: PromptStrategy = None, history: Prompt = None, handlers=(),
                  context: PromptContext = None, rest_api: str = "",
-                 schemas: dict = None):
+                 schemas: dict = None, endpoints: dict = None,):
         """
         Initializes the PromptEngineer with a specific strategy and handlers for LLM and responses.
 
@@ -27,7 +28,7 @@ def __init__(self, strategy: PromptStrategy = None, history: Prompt = None, hand
         self.strategy = strategy
         self.rest_api = rest_api
         self.llm_handler, self.response_handler = handlers
-        self.prompt_helper = PromptGenerationHelper(response_handler=self.response_handler, schemas=schemas or {})
+        self.prompt_helper = PromptGenerationHelper(response_handler=self.response_handler, schemas=schemas or {}, endpoints=endpoints)
         self.context = context
         self.turn = 0
         self._prompt_history = history or []
@@ -42,7 +43,7 @@ def __init__(self, strategy: PromptStrategy = None, history: Prompt = None, hand
                                                                    self.turn: {"content": "initial_prompt"}})
         }
 
-        self.purpose = None
+        self.purpose =  PromptPurpose.AUTHENTICATION_AUTHORIZATION
 
     def generate_prompt(self, turn:int, move_type="explore", hint=""):
         """

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

@@ -15,7 +15,7 @@ class PromptGenerationHelper(object):
         schemas (dict): A dictionary of schemas used for constructing HTTP requests.
     """
 
-    def __init__(self, response_handler:ResponseHandler=None, schemas:dict={}):
+    def __init__(self, response_handler: ResponseHandler = None, schemas: dict = {}, endpoints: dict = {}):
         """
         Initializes the PromptAssistant with a response handler and downloads necessary NLTK models.
 
@@ -28,13 +28,7 @@ def __init__(self, response_handler:ResponseHandler=None, schemas:dict={}):
         self.endpoint_methods = {}
         self.endpoint_found_methods = {}
         self.schemas = schemas
-
-        # Download NLTK models if not already installed
-        nltk.download('punkt')
-        nltk.download('stopwords')
-
-
-
+        self.endpoints = endpoints
 
     def get_endpoints_needing_help(self):
         """
@@ -106,6 +100,8 @@ def token_count(self, text):
         Returns:
             int: The number of tokens in the input text.
         """
+        if not isinstance(text, str):
+            text = str(text)
         tokens = re.findall(r'\b\w+\b', text)
         words = [token.strip("'") for token in tokens if token.strip("'").isalnum()]
         return len(words)
@@ -124,6 +120,7 @@ def check_prompt(self, previous_prompt: list, steps: str, max_tokens: int = 900)
         """
 
         def validate_prompt(prompt):
+            print(f'Prompt: {prompt}')
             if self.token_count(prompt) <= max_tokens:
                 return prompt
             shortened_prompt = self.response_handler.get_response_for_prompt("Shorten this prompt: " + prompt)
@@ -135,7 +132,7 @@ def validate_prompt(prompt):
             if isinstance(steps, list):
                 potential_prompt = "\n".join(str(element) for element in steps)
             else:
-                potential_prompt = str(steps) +"\n"
+                potential_prompt = str(steps) + "\n"
             return validate_prompt(potential_prompt)
 
         return validate_prompt(previous_prompt)

src/hackingBuddyGPT/usecases/web_api_testing/prompt_generation/prompts/basic_prompt.py

@@ -41,7 +41,7 @@ def __init__(self, context: PromptContext = None, planning_type: PlanningType =
         self.pentesting_information: Optional[PenTestingInformation] = None
 
         if self.context == PromptContext.PENTESTING:
-            self.pentesting_information = PenTestingInformation(schemas=prompt_helper.schemas)
+            self.pentesting_information = PenTestingInformation(schemas=prompt_helper.schemas, endpoints=prompt_helper.endpoints)
 
     @abstractmethod
     def generate_prompt(self, move_type: str, hint: Optional[str], previous_prompt: Optional[str],

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

@@ -1,7 +1,9 @@
 from typing import List, Optional
 
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy, PromptContext, PromptPurpose
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.task_planning.task_planning_prompt import TaskPlanningPrompt
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy, \
+    PromptContext, PromptPurpose
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.task_planning.task_planning_prompt import \
+    TaskPlanningPrompt
 
 
 class ChainOfThoughtPrompt(TaskPlanningPrompt):
@@ -120,20 +122,23 @@ def _get_pentesting_steps(self, move_type: str) -> List[str]:
             List[str]: A list of steps for the chain-of-thought strategy in the pentesting context.
         """
         if move_type == "explore":
-            purpose = list(self.pentesting_information.explore_steps.keys())[0]
-            step = self.pentesting_information.explore_steps[purpose]
-            if step not in self.explored_steps:
-                if len(step) > 1:
-                    step = self.pentesting_information.explore_steps[purpose][0]
-                    if len(self.pentesting_information.explore_steps[purpose]) == 0:
+            if len(self.pentesting_information.explore_steps.keys()) > 0:
+                purpose = list(self.pentesting_information.explore_steps.keys())[0]
+                step = self.pentesting_information.explore_steps[purpose]
+                if step not in self.explored_steps:
+                    if len(step) > 1:
+                        step = self.pentesting_information.explore_steps[purpose][0]
+                        # Delete the first item from the list, automatically shifting the remaining items up
                         del self.pentesting_information.explore_steps[purpose][0]
-                prompt = step
-                self.purpose = purpose
-                self.explored_steps.append(step)
-                if len(step) == 1:
-                    del self.pentesting_information.explore_steps[purpose]
-
-                print(f'prompt: {prompt}')
-                return prompt
+                    prompt = step
+                    self.purpose = purpose
+                    self.explored_steps.append(step)
+                    if len(step) == 1:
+                        del self.pentesting_information.explore_steps[purpose]
+
+                    print(f'prompt: {prompt}')
+                    return prompt
+            else:
+                return ""
         else:
             return ["Look for exploits."]

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

@@ -77,7 +77,7 @@ def analyze_response(self, raw_response: str, prompt_history: list) -> tuple[dic
             for step in steps:
                 prompt_history, response = self.process_step(step, prompt_history)
                 llm_responses.append(response)
-                print(f'Response:{response}')
+                #print(f'Response:{response}')
 
         return llm_responses
 
@@ -98,13 +98,13 @@ def parse_http_response(self, raw_response: str):
 
         match = re.match(r"HTTP/1\.1 (\d{3}) (.*)", status_line)
         status_code = int(match.group(1)) if match else None
-        if body.__contains__("<html"):
+        if body.__contains__("<!DOCTYPE"):
             body = ""
 
         elif status_code in [500, 400, 404, 422]:
             body = body
         else:
-            print(f'Body:{body}')
+            #print(f'Body:{body}')
             if body != '' or body != "":
                 body = json.loads(body)
             if isinstance(body, list) and len(body) > 1:

src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py

@@ -9,7 +9,8 @@
 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.prompt_generation.information.prompt_information import PromptContext
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptContext, \
+    PromptPurpose
 from hackingBuddyGPT.usecases.web_api_testing.utils.custom_datatypes import Prompt, Context
 from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing import OpenAPISpecificationParser
 from hackingBuddyGPT.usecases.web_api_testing.documentation.report_handler import ReportHandler
@@ -23,7 +24,7 @@
 
 
 # OpenAPI specification file path
-openapi_spec_filename = "/home/diana/Desktop/masterthesis/00/hackingBuddyGPT/src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_spec/openapi_spec_2024-08-16_14-14-07.yaml"
+openapi_spec_filename = "src/hackingBuddyGPT/usecases/web_api_testing/documentation/openapi_spec/openapi_spec_2024-09-03_10-22-09.yaml"
 
 
 class SimpleWebAPITesting(Agent):
@@ -76,6 +77,7 @@ def init(self) -> None:
         self._response_handler: ResponseHandler = ResponseHandler(self._llm_handler)
         self._report_handler: ReportHandler = ReportHandler()
         self._setup_initial_prompt()
+        self.purpose =  PromptPurpose.AUTHENTICATION_AUTHORIZATION
 
     def _setup_initial_prompt(self) -> None:
         """
@@ -96,13 +98,16 @@ def _setup_initial_prompt(self) -> None:
         handlers = (self._llm_handler, self._response_handler)
         schemas: Dict[str, Any] = self._openapi_specification["components"]["schemas"] if os.path.exists(
             openapi_spec_filename) else {}
+        endpoints: Dict[str, Any] = self._openapi_specification["paths"].keys() if os.path.exists(
+            openapi_spec_filename) else {}
         self.prompt_engineer: PromptEngineer = PromptEngineer(
             strategy=PromptStrategy.CHAIN_OF_THOUGHT,
             history=self._prompt_history,
             handlers=handlers,
             context=PromptContext.PENTESTING,
             rest_api=self.host,
-            schemas=schemas
+            schemas=schemas,
+            endpoints= endpoints
         )
 
     def all_http_methods_found(self) -> None:
@@ -136,11 +141,19 @@ def perform_round(self, turn: int) -> None:
         Args:
             turn (int): The current round number.
         """
-        prompt = self.prompt_engineer.generate_prompt(turn)
+        self._perform_prompt_generation(turn)
+    def _perform_prompt_generation(self, turn: int) -> None:
         response: Any
         completion: Any
-        response, completion = self._llm_handler.call_llm(prompt)
-        self._handle_response(completion, response, self.prompt_engineer.purpose)
+        while self.purpose == self.prompt_engineer.purpose:
+            print(f'Self purpose: {self.purpose}')
+            print(f'prompt engineer purpose: {self.purpose}')
+            prompt = self.prompt_engineer.generate_prompt(turn)
+            response, completion = self._llm_handler.call_llm(prompt)
+            self._handle_response(completion, response, self.prompt_engineer.purpose)
+        print(f'Self purpose: {self.purpose}')
+        print(f'prompt engineer purpose: {self.purpose}')
+        self.purpose = self.prompt_engineer.purpose
 
     def _handle_response(self, completion: Any, response: Any, purpose: str) -> None:
         """
@@ -173,6 +186,8 @@ def _handle_response(self, completion: Any, response: Any, purpose: str) -> None
         self.all_http_methods_found()
 
 
+
+
 @use_case("Minimal implementation of a web API testing use case")
 class SimpleWebAPITestingUseCase(AutonomousAgentUseCase[SimpleWebAPITesting]):
     """