Merge pull request #122 from ipa-lab/adapt_prompt_engineer

Adapt prompt engineer

Author: andreashappe

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):