Adapt prompt engineer

src/hackingBuddyGPT/usecases/web_api_testing/__init__.py

@@ -2,5 +2,4 @@
 from .simple_web_api_testing import SimpleWebAPITesting
 from . import response_processing
 from . import documentation
-from . import prompt_generation
 from . import testing

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

@@ -5,7 +5,7 @@
 import yaml
 from hackingBuddyGPT.capabilities.yamlFile import YAMLFile
 from hackingBuddyGPT.usecases.web_api_testing.documentation.pattern_matcher import PatternMatcher
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information import PromptStrategy
+from hackingBuddyGPT.utils.prompt_generation.information import PromptStrategy
 from hackingBuddyGPT.usecases.web_api_testing.response_processing import ResponseHandler
 from hackingBuddyGPT.usecases.web_api_testing.utils import LLMHandler
 
@@ -42,7 +42,8 @@ def __init__(self, llm_handler: LLMHandler, response_handler: ResponseHandler, s
         self.query_params = {}
         self.endpoint_methods = {}
         self.endpoint_examples = {}
-        self.filename = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.yaml"
+        date = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+        self.filename = f"{name}_spec.yaml"
         self.openapi_spec = {
             "openapi": "3.0.0",
             "info": {
@@ -57,7 +58,7 @@ def __init__(self, llm_handler: LLMHandler, response_handler: ResponseHandler, s
         self.llm_handler = llm_handler
         current_path = os.path.dirname(os.path.abspath(__file__))
 
-        self.file_path = os.path.join(current_path, "openapi_spec", str(strategy).split(".")[1].lower(), name.lower())
+        self.file_path = os.path.join(current_path, "openapi_spec", str(strategy).split(".")[1].lower(), name.lower(), date)
         os.makedirs(self.file_path, exist_ok=True)
         self.file = os.path.join(self.file_path, self.filename)
 

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

@@ -2,7 +2,7 @@
 import re
 from typing import Any, Dict, Optional, Tuple
 
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptPurpose
+from hackingBuddyGPT.utils.prompt_generation.information import PromptPurpose
 
 
 class ResponseAnalyzer:

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

@@ -1,13 +1,13 @@
 import json
 import re
-from typing import Any, Dict, Tuple, List
+from typing import Any, Dict
 from unittest.mock import MagicMock
 
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information import (
+from hackingBuddyGPT.utils.prompt_generation.information import (
     PenTestingInformation,
 )
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import (
+from hackingBuddyGPT.utils.prompt_generation.information import (
     PromptPurpose,
 )
 from hackingBuddyGPT.usecases.web_api_testing.utils import LLMHandler

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

@@ -1,6 +1,5 @@
 import copy
 import json
-import os.path
 import re
 from collections import Counter
 from itertools import cycle
@@ -12,9 +11,9 @@
 from rich.panel import Panel
 
 from hackingBuddyGPT.usecases.web_api_testing.documentation.pattern_matcher import PatternMatcher
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation import PromptGenerationHelper
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information import PromptContext
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.pentesting_information import (
+from hackingBuddyGPT.utils.prompt_generation import PromptGenerationHelper
+from hackingBuddyGPT.utils.prompt_generation.information import PromptContext
+from hackingBuddyGPT.utils.prompt_generation.information import (
     PenTestingInformation,
 )
 from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_analyzer_with_llm import (
@@ -510,6 +509,8 @@ def handle_http_response(self, response: Any, prompt_history: Any, log: Any, com
             self.last_path = request_path
 
             status_message = self.check_if_successful(is_successful, request_path, result_dict, result_str, categorized_endpoints)
+            log.console.print(Panel(status_message, title="system"))
+
             prompt_history.append(tool_message(status_message, tool_call_id))
 
         else:
@@ -777,9 +778,6 @@ def update_step_and_category():
         elif self.prompt_helper.current_step == 7 and not self.prompt_helper._get_root_level_endpoints(self.name):
             update_step_and_category()
 
-    import random
-    from urllib.parse import urlencode
-
     def create_common_query_for_endpoint(self, endpoint):
         """
         Constructs complete URLs with one query parameter for each API endpoint.
@@ -948,6 +946,8 @@ def check_if_successful(self, is_successful, request_path, result_dict, result_s
             error_msg = result_dict.get("error", {}).get("message", "unknown error") if isinstance(
                 result_dict.get("error", {}), dict) else result_dict.get("error", "unknown error")
             self.no_new_endpoint_counter +=1
+            if error_msg == "unknown error" and (result_str.startswith("4") or result_str.startswith("5")):
+                error_msg = result_str
 
             if result_str.startswith("400") or result_str.startswith("401") or result_str.startswith("403"):
                 status_message = f"{request_path} is a correct endpoint, but encountered an error: {error_msg}"
@@ -960,7 +960,6 @@ def check_if_successful(self, is_successful, request_path, result_dict, result_s
                 if error_msg not in self.prompt_helper.correct_endpoint_but_some_error:
                     self.prompt_helper.correct_endpoint_but_some_error[error_msg] = []
                 self.prompt_helper.correct_endpoint_but_some_error[error_msg].append(request_path)
-                self.prompt_helper.hint_for_next_round = error_msg
             else:
                 self.prompt_helper.unsuccessful_paths.append(request_path)
                 status_message = f"{request_path} is not a correct endpoint; Reason: {error_msg}"

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -2,16 +2,18 @@
 from dataclasses import field
 from typing import Dict
 
+from rich.panel import Panel
+
 from hackingBuddyGPT.capabilities import Capability
 from hackingBuddyGPT.capabilities.http_request import HTTPRequest
 from hackingBuddyGPT.capabilities.record_note import RecordNote
 from hackingBuddyGPT.usecases.agents import Agent
 from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case
 from hackingBuddyGPT.usecases.web_api_testing.documentation.openapi_specification_handler import \
     OpenAPISpecificationHandler
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation import PromptGenerationHelper
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptContext
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompt_engineer import PromptEngineer
+from hackingBuddyGPT.utils.prompt_generation import PromptGenerationHelper
+from hackingBuddyGPT.utils.prompt_generation.information import PromptContext
+from hackingBuddyGPT.utils.prompt_generation.prompt_engineer import PromptEngineer
 from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_handler import ResponseHandler
 from hackingBuddyGPT.usecases.web_api_testing.utils import LLMHandler
 from hackingBuddyGPT.usecases.web_api_testing.utils.configuration_handler import ConfigurationHandler
@@ -54,6 +56,11 @@ class SimpleWebAPIDocumentation(Agent):
         default="",
     )
 
+    prompt_file: str = parameter(
+        desc="prompt file name",
+        default="",
+    )
+
 
     _http_method_description: str = parameter(
         desc="Pattern description for expected HTTP methods in the API response",
@@ -155,10 +162,11 @@ def _initialize_handlers(self, config, description, token, name, initial_prompt)
 
         self._prompt_engineer = PromptEngineer(
             strategy=self.strategy,
-            context=self._prompt_context,
+            context=PromptContext.DOCUMENTATION,
             prompt_helper=self.prompt_helper,
             open_api_spec=self._documentation_handler.openapi_spec,
-            rest_api_info=(token, self.host, self._correct_endpoints, self.categorized_endpoints)
+            rest_api_info=(token, self.host, self._correct_endpoints, self.categorized_endpoints),
+            prompt_file=self.prompt_file
         )
         self._evaluator = Evaluator(config=config)
 
@@ -376,6 +384,8 @@ def run_documentation(self, turn: int, move_type: str) -> None:
             prompt = self._prompt_engineer.generate_prompt(turn=turn, move_type=move_type,
                                                            prompt_history=self._prompt_history)
             response, completion = self._llm_handler.execute_prompt_with_specific_capability(prompt,"http_request" )
+            self.log.console.print(Panel(prompt[-1]["content"], title="system"))
+
             is_good, self._prompt_history, result, result_str = self._response_handler.handle_response(response,
                                                                                                        completion,
                                                                                                        self._prompt_history,

src/hackingBuddyGPT/usecases/web_api_testing/simple_web_api_testing.py

@@ -15,14 +15,13 @@
 from hackingBuddyGPT.capabilities.record_note import RecordNote
 from hackingBuddyGPT.usecases.agents import Agent
 from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation import PromptGenerationHelper
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information import PenTestingInformation
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptContext, \
-    PromptPurpose
+from hackingBuddyGPT.utils.prompt_generation import PromptGenerationHelper
+from hackingBuddyGPT.utils.prompt_generation.information import PenTestingInformation
+from hackingBuddyGPT.utils.prompt_generation.information import PromptPurpose
 from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing import OpenAPISpecificationParser
 from hackingBuddyGPT.usecases.web_api_testing.documentation.report_handler import ReportHandler
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptContext
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompt_engineer import PromptEngineer, PromptStrategy
+from hackingBuddyGPT.utils.prompt_generation.information import PromptContext
+from hackingBuddyGPT.utils.prompt_generation.prompt_engineer import PromptEngineer
 from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_analyzer_with_llm import \
     ResponseAnalyzerWithLLM
 from hackingBuddyGPT.usecases.web_api_testing.response_processing.response_handler import ResponseHandler

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

@@ -1,7 +1,7 @@
 import json
 import os
 
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information import PromptStrategy, PromptContext
+from hackingBuddyGPT.utils.prompt_generation.information import PromptStrategy
 
 
 class ConfigurationHandler(object):

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

@@ -10,7 +10,7 @@
 import pandas
 
 from hackingBuddyGPT.usecases.web_api_testing.documentation.parsing import OpenAPISpecificationParser
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import (
+from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import (
     PromptPurpose,
 )
 from faker import Faker

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

@@ -1,16 +1,15 @@
-from itertools import cycle
+from typing import Any
 
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import (
+from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import (
     PromptContext,
-    PromptStrategy, PromptPurpose,
-)
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompt_generation_helper import (
+    PromptStrategy, )
+from hackingBuddyGPT.utils.prompt_generation.prompt_generation_helper import (
     PromptGenerationHelper,
 )
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.state_learning import (
+from hackingBuddyGPT.utils.prompt_generation.prompts.state_learning import (
     InContextLearningPrompt,
 )
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.task_planning import (
+from hackingBuddyGPT.utils.prompt_generation.prompts.task_planning import (
     ChainOfThoughtPrompt,
     TreeOfThoughtPrompt,
 )
@@ -35,6 +34,7 @@ def __init__(
             open_api_spec: dict = None,
             prompt_helper: PromptGenerationHelper = None,
             rest_api_info: tuple = None,
+            prompt_file : Any = None
     ):
 
         """
@@ -58,16 +58,17 @@ def __init__(
 
         strategies = {
             PromptStrategy.CHAIN_OF_THOUGHT: ChainOfThoughtPrompt(
-                context=context, prompt_helper=self.prompt_helper,
+                context=context, prompt_helper=self.prompt_helper, prompt_file = prompt_file
             ),
             PromptStrategy.TREE_OF_THOUGHT: TreeOfThoughtPrompt(
-                context=context, prompt_helper=self.prompt_helper
+                context=context, prompt_helper=self.prompt_helper,  prompt_file = prompt_file
             ),
             PromptStrategy.IN_CONTEXT: InContextLearningPrompt(
                 context=context,
                 prompt_helper=self.prompt_helper,
                 context_information={self.turn: {"content": "initial_prompt"}},
-                open_api_spec=open_api_spec
+                open_api_spec=open_api_spec,
+                prompt_file=prompt_file
             ),
         }
 

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

@@ -1,9 +1,10 @@
+import os.path
 from abc import ABC, abstractmethod
-from typing import Optional
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information import (
+from typing import Optional, Any
+from hackingBuddyGPT.utils.prompt_generation.information import (
     PenTestingInformation,
 )
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import (
+from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import (
     PlanningType,
     PromptContext,
     PromptStrategy, PromptPurpose,
@@ -31,6 +32,7 @@ def __init__(
             planning_type: PlanningType = None,
             prompt_helper=None,
             strategy: PromptStrategy = None,
+            prompt_file: Any =None
     ):
         """
         Initializes the BasicPrompt with a specific context, prompt helper, and strategy.
@@ -44,6 +46,9 @@ def __init__(
         self.transformed_steps = {}
         self.open_api_spec = {}
         self.context = context
+        if context is None:
+            if os.path.exists(prompt_file):
+                self.prompt_file = prompt_file
         self.planning_type = planning_type
         self.prompt_helper = prompt_helper
         self.strategy = strategy
@@ -79,27 +84,15 @@ def generate_prompt(
         pass
 
     def get_documentation_steps(self):
-        """
-        Returns a predefined list of endpoint exploration steps based on the target API host.
-
-        These steps are used to guide automated documentation of a web API by progressively
-        discovering and querying endpoints using GET requests. The process follows a structured
-        hierarchy from root-level endpoints to more complex nested endpoints and those with query parameters.
-
-        Returns:
-            List[List[str]]: A list of steps, each step being a list of instruction strings.
-        """
-
-        # Define specific documentation steps based on the given strategy
-
         return [
-            [f"Objective: Identify all accessible endpoints via GET requests for {self.prompt_helper.host}. {self.prompt_helper._description}"],
+            [
+                f"Objective: Identify all accessible endpoints via GET requests for {self.prompt_helper.host}. {self.prompt_helper._description}"],
             [
                 f""" Query root-level resource endpoints.
-                               Find root-level endpoints for {self.prompt_helper.host}.
-                               Only send GET requests to root-level endpoints with a single path component after the root. This means each path should have exactly one '/' followed by a single word (e.g., '/users', '/products').  
-                               1. Send GET requests to new paths only, avoiding any in the lists above.
-                               2. Do not reuse previously tested paths."""
+                                      Find root-level endpoints for {self.prompt_helper.host}.
+                                      Only send GET requests to root-level endpoints with a single path component after the root. This means each path should have exactly one '/' followed by a single word (e.g., '/users', '/products').  
+                                      1. Send GET requests to new paths only, avoiding any in the lists above.
+                                      2. Do not reuse previously tested paths."""
 
             ],
             [
@@ -114,7 +107,6 @@ def get_documentation_steps(self):
                 "Identify subresource endpoints of the form `/resource/other_resource`.",
                 "Query these endpoints to check if they return data related to the main resource without requiring an `id` parameter."
 
-
             ],
 
             [
@@ -133,6 +125,7 @@ def get_documentation_steps(self):
                 "Construct and make GET requests to these endpoints using common query parameters (e.g. `/resource?param1=1&param2=3`) or based on documentation hints, testing until a valid request with query parameters is achieved."
             ]
         ]
+
     def extract_properties(self):
         """
            Extracts example values and data types from the 'Post' schema in the OpenAPI specification.
@@ -176,6 +169,19 @@ def sort_previous_prompt(self, previous_prompt):
             sorted_list.append(previous_prompt[i])
         return sorted_list
 
+    def parse_prompt_file(self):
+        with open(self.prompt_file, "r", encoding="utf-8") as f:
+            content = f.read()
+        blocks = content.strip().split('---')
+        prompt_blocks = []
+
+        for block in blocks:
+            block = block.replace("{host}", self.prompt_helper.host).replace("{description}", self.prompt_helper._description)
+            lines = [line.strip() for line in block.strip().splitlines() if line.strip()]
+            if lines:
+                prompt_blocks.append(lines)
+
+        return prompt_blocks
 
     def extract_endpoints_from_prompts(self, step):
         """