From db79b7412603d9fd365229e3a128079499e5f8b9 Mon Sep 17 00:00:00 2001 From: Grey_D Date: Sat, 9 Nov 2024 14:57:10 +0800 Subject: [PATCH 1/2] update all the params --- pentestgpt/_version.py | 2 +- pentestgpt/main.py | 56 +++++++----- pentestgpt/test_connection.py | 114 +++++++++++------------- pentestgpt/utils/APIs/gemini_api.py | 7 +- pentestgpt/utils/APIs/module_import.py | 2 +- pentestgpt/utils/pentest_gpt.py | 9 +- pentestgpt/utils/pentest_gpt_rebuilt.py | 9 +- pentestgpt/utils/prompt_select.py | 3 +- pyproject.toml | 2 +- setup.py | 2 +- 10 files changed, 97 insertions(+), 109 deletions(-) diff --git a/pentestgpt/_version.py b/pentestgpt/_version.py index b80d381..83412d9 100644 --- a/pentestgpt/_version.py +++ b/pentestgpt/_version.py @@ -1 +1 @@ -__version__ = '"0.14.0"' +__version__ = '""' diff --git a/pentestgpt/main.py b/pentestgpt/main.py index a527974..93ab2aa 100644 --- a/pentestgpt/main.py +++ b/pentestgpt/main.py @@ -1,11 +1,14 @@ -from dataclasses import dataclass +import argparse import sys +from dataclasses import dataclass from typing import Optional -import argparse + from loguru import logger -from pentestgpt.test_connection import main as test_connection + +from pentestgpt.test_connection import test_connection from pentestgpt.utils.pentest_gpt import pentestGPT + @dataclass class PentestConfig: log_dir: str @@ -14,17 +17,18 @@ class PentestConfig: use_logging: bool use_api: bool + class PentestGPTCLI: - + DEFAULT_CONFIG = { "log_dir": "logs", - "reasoning_model": "gpt-4-o", - "parsing_model": "gpt-4-o", + "reasoning_model": "gpt-4o", + "parsing_model": "gpt-4o", } - + VALID_MODELS = { - "reasoning": ["gpt-4", "gpt-4-turbo"], - "parsing": ["gpt-4-turbo", "gpt-3.5-turbo-16k"] + "reasoning": ["gpt-4", "gpt-4-turbo", "gpt-4o", "gpt-3.5-turbo-16k"], + "parsing": ["gpt-4-turbo", "gpt-3.5-turbo-16k", "gpt-4o"], } def __init__(self): @@ -32,44 +36,44 @@ def __init__(self): def _create_parser(self) -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="PentestGPT") - + parser.add_argument( "--log_dir", type=str, default=self.DEFAULT_CONFIG["log_dir"], - help="Path to the log directory for storing conversations" + help="Path to the log directory for storing conversations", ) - + parser.add_argument( "--reasoning_model", type=str, default=self.DEFAULT_CONFIG["reasoning_model"], choices=self.VALID_MODELS["reasoning"], - help="Model for higher-level cognitive tasks" + help="Model for higher-level cognitive tasks", ) - + parser.add_argument( "--parsing_model", type=str, default=self.DEFAULT_CONFIG["parsing_model"], choices=self.VALID_MODELS["parsing"], - help="Model for structural and grammatical language processing" + help="Model for structural and grammatical language processing", ) - + parser.add_argument( "--logging", action="store_true", default=False, - help="Enable data collection through langfuse logging" + help="Enable data collection through langfuse logging", ) - + parser.add_argument( "--useAPI", action="store_true", default=True, - help="Deprecated: Set to False only for testing with cookie" + help="Deprecated: Set to False only for testing with cookie", ) - + return parser def parse_args(self) -> PentestConfig: @@ -79,9 +83,10 @@ def parse_args(self) -> PentestConfig: reasoning_model=args.reasoning_model, parsing_model=args.parsing_model, use_logging=args.logging, - use_api=args.useAPI + use_api=args.useAPI, ) + def check_connection() -> bool: try: return test_connection() @@ -89,6 +94,7 @@ def check_connection() -> bool: logger.error(f"Connection test failed: {e}") return False + def run_pentest(config: PentestConfig) -> None: try: pentest_handler = pentestGPT( @@ -96,22 +102,24 @@ def run_pentest(config: PentestConfig) -> None: parsing_model=config.parsing_model, useAPI=config.use_api, log_dir=config.log_dir, - use_langfuse_logging=config.use_logging + use_langfuse_logging=config.use_logging, ) pentest_handler.main() except Exception as e: logger.error(f"PentestGPT execution failed: {e}") sys.exit(1) + def main(): cli = PentestGPTCLI() config = cli.parse_args() - + if not check_connection(): logger.error("Connection test failed. Exiting...") sys.exit(1) - + run_pentest(config) + if __name__ == "__main__": main() diff --git a/pentestgpt/test_connection.py b/pentestgpt/test_connection.py index 35a7cec..443a6f7 100644 --- a/pentestgpt/test_connection.py +++ b/pentestgpt/test_connection.py @@ -1,5 +1,4 @@ import argparse -from pathlib import Path import loguru import openai @@ -8,7 +7,6 @@ from pentestgpt._version import __version__ from pentestgpt.config.chat_config import ChatGPTConfig from pentestgpt.utils.APIs.chatgpt_api import ChatGPTAPI -from pentestgpt.utils.chatgpt import ChatGPT logger = loguru.logger @@ -18,78 +16,68 @@ def get_project_version(): return __version__ +def test_connection(log_dir="logs", base_url=ChatGPTConfig().api_base): + """ + Test connection function, can be called directly or from pentestgpt. + + Args: + log_dir (str): Directory for logging connection tests. + base_url (str): Base URL for the OpenAI API. + """ + # Set up logging + logger.add(f"{log_dir}/chatgpt_connection_test.log", level="ERROR") + chatgpt_config = ChatGPTConfig(api_base=base_url) + console = Console() + + # Print version + version = get_project_version() + console.print( + f"Testing the connection for PentestGPT v{version}", style="bold green" + ) + + # Test connection with GPT-3.5 and GPT-4o + models_to_test = ["gpt-3.5-turbo-16k", "gpt-4o"] + for model in models_to_test: + print(f"#### Testing connection for OpenAI API with model {model}") + try: + chatgpt_config.model = model + chatgpt = ChatGPTAPI(chatgpt_config) + openai.api_key = chatgpt_config.openai_key + chatgpt.send_new_message("Hi, how are you?") + console.print( + f"Connection successful with model {model}. Use to start.", + style="bold green", + ) + except Exception as e: + logger.error(e) + console.print( + f"Failed to connect with model {model}. Check API key and configuration.", + style="bold red", + ) + print("Error:", e) + return False + return True + + def main(): - parser = argparse.ArgumentParser(description="PentestGPTTestConnection") + # Parse arguments only if running this script directly + parser = argparse.ArgumentParser(description="PentestGPT Test Connection") parser.add_argument( - "--logDir", + "--log_dir", type=str, default="logs", - help="Log file directory for PentestGPTTestConnection", + help="Directory for logging connection tests", ) - parser.add_argument( - "--baseUrl", + "--base_url", type=str, default=ChatGPTConfig().api_base, - help="Base URL for OpenAI API,default: https://api.openai.com/v1", + help="Base URL for OpenAI API, default: https://api.openai.com/v1", ) - args = parser.parse_args() - logger.add(args.logDir + "/chatgpt_connection_test.log", level="ERROR") - - chatgpt_config = ChatGPTConfig(api_base=args.baseUrl) - console = Console() - - # print version - version = get_project_version() - console.print( - f"You're testing the connection for PentestGPT v{version}", style="bold green" - ) - - # successful connection bool - can_connect = False - - # 1. test the connection for chatgpt api with GPT-3.5 - print("#### Test connection for OpenAI api (GPT-3.5)") - try: - chatgpt_config.model = "gpt-3.5-turbo-16k" - chatgpt = ChatGPTAPI(chatgpt_config) - openai.api_key = chatgpt_config.openai_key - result, conversation_id = chatgpt.send_new_message("Hi how are you?") - console.print( - "1. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use ", - style="bold green", - ) - can_connect = True - except Exception as e: # use a general exception first. Update later for debug - logger.error(e) - console.print( - "1. The OpenAI API key is not properly configured. The likely reason is that you do not link a payment method to OpenAI so your key is not active. \nPlease follow README to update OpenAI API key through `export OPENAI_API_KEY=<>`", - style="bold red", - ) - print("The error is below:", e) - - # 2. test the connection for chatgpt api with GPT-4 - print("#### Test connection for OpenAI api (GPT-4)") - try: - chatgpt_config.model = "gpt-4o" - chatgpt = ChatGPTAPI(chatgpt_config) - openai.api_key = chatgpt_config.openai_key - result, conversation_id = chatgpt.send_new_message("Hi how are you?") - console.print( - "1. You're connected with OpenAI API. You have GPT-4(o) access. To start PentestGPT, please use ", - style="bold green", - ) - can_connect = True - except Exception as e: # use a general exception first. Update later for debug - logger.error(e) - console.print( - "2. The OpenAI API key is not properly configured. Please check the error below:", - style="bold red", - ) - print("The error is below:", e) - return can_connect + # Call test_connection with arguments from CLI + test_connection(log_dir=args.log_dir, base_url=args.base_url) if __name__ == "__main__": diff --git a/pentestgpt/utils/APIs/gemini_api.py b/pentestgpt/utils/APIs/gemini_api.py index 46f4817..c8cab8b 100644 --- a/pentestgpt/utils/APIs/gemini_api.py +++ b/pentestgpt/utils/APIs/gemini_api.py @@ -10,11 +10,8 @@ import google.generativeai as genai import loguru import tiktoken -from google.generativeai.types import ( - HarmBlockThreshold, - HarmCategory, - SafetySettingDict, -) +from google.generativeai.types import (HarmBlockThreshold, HarmCategory, + SafetySettingDict) from langfuse.model import InitialGeneration, Usage from tenacity import * diff --git a/pentestgpt/utils/APIs/module_import.py b/pentestgpt/utils/APIs/module_import.py index 1a2697f..a0d7254 100644 --- a/pentestgpt/utils/APIs/module_import.py +++ b/pentestgpt/utils/APIs/module_import.py @@ -14,7 +14,7 @@ "module_name": "chatgpt_api", "class_name": "ChatGPTAPI", }, - "gpt-4-o": { + "gpt-4o": { "config_name": "GPT4O", "module_name": "chatgpt_api", "class_name": "ChatGPTAPI", diff --git a/pentestgpt/utils/pentest_gpt.py b/pentestgpt/utils/pentest_gpt.py index 8acd750..34ec545 100644 --- a/pentestgpt/utils/pentest_gpt.py +++ b/pentestgpt/utils/pentest_gpt.py @@ -17,12 +17,9 @@ from pentestgpt.utils.APIs.module_import import dynamic_import from pentestgpt.utils.chatgpt import ChatGPT from pentestgpt.utils.prompt_select import prompt_ask, prompt_select -from pentestgpt.utils.task_handler import ( - local_task_entry, - localTaskCompleter, - main_task_entry, - mainTaskCompleter, -) +from pentestgpt.utils.task_handler import (local_task_entry, + localTaskCompleter, main_task_entry, + mainTaskCompleter) from pentestgpt.utils.web_parser import google_search logger = loguru.logger diff --git a/pentestgpt/utils/pentest_gpt_rebuilt.py b/pentestgpt/utils/pentest_gpt_rebuilt.py index 1459045..785526a 100644 --- a/pentestgpt/utils/pentest_gpt_rebuilt.py +++ b/pentestgpt/utils/pentest_gpt_rebuilt.py @@ -17,12 +17,9 @@ from pentestgpt.utils.APIs.module_import import dynamic_import from pentestgpt.utils.chatgpt import ChatGPT from pentestgpt.utils.prompt_select import prompt_ask, prompt_select -from pentestgpt.utils.task_handler import ( - local_task_entry, - localTaskCompleter, - main_task_entry, - mainTaskCompleter, -) +from pentestgpt.utils.task_handler import (local_task_entry, + localTaskCompleter, main_task_entry, + mainTaskCompleter) from pentestgpt.utils.web_parser import google_search logger = loguru.logger diff --git a/pentestgpt/utils/prompt_select.py b/pentestgpt/utils/prompt_select.py index 94ba3e4..7fee7d4 100644 --- a/pentestgpt/utils/prompt_select.py +++ b/pentestgpt/utils/prompt_select.py @@ -3,7 +3,8 @@ from prompt_toolkit.application import Application from prompt_toolkit.formatted_text import HTML from prompt_toolkit.key_binding.defaults import load_key_bindings -from prompt_toolkit.key_binding.key_bindings import KeyBindings, merge_key_bindings +from prompt_toolkit.key_binding.key_bindings import (KeyBindings, + merge_key_bindings) from prompt_toolkit.layout import Layout from prompt_toolkit.layout.containers import HSplit from prompt_toolkit.shortcuts import prompt diff --git a/pyproject.toml b/pyproject.toml index 2c8a1ad..b1844ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pentestgpt" -version = "0.14.0" +version = "0.14.1" description = "PentestGPT is an LLM-powered penetration testing tool." authors = ["Gelei Deng "] license = "MIT" diff --git a/setup.py b/setup.py index 5524123..1779e92 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ setup( name="pentestgpt", - version="0.14.0", + version="", description="PentestGPT, a GPT-empowered penetration testing tool", long_description=""" PentestGPT is a penetration testing tool empowered by ChatGPT. From b831d4e9d0e4471cf0ec313f62927a07329483a8 Mon Sep 17 00:00:00 2001 From: "deepsource-autofix[bot]" <62050782+deepsource-autofix[bot]@users.noreply.github.com> Date: Sat, 9 Nov 2024 06:57:45 +0000 Subject: [PATCH 2/2] style: format code with Black This commit fixes the style issues introduced in db79b74 according to the output from Black. Details: https://github.com/GreyDGL/PentestGPT/pull/251 --- pentestgpt/utils/APIs/gemini_api.py | 7 +++++-- pentestgpt/utils/pentest_gpt.py | 9 ++++++--- pentestgpt/utils/pentest_gpt_rebuilt.py | 9 ++++++--- pentestgpt/utils/prompt_select.py | 3 +-- 4 files changed, 18 insertions(+), 10 deletions(-) diff --git a/pentestgpt/utils/APIs/gemini_api.py b/pentestgpt/utils/APIs/gemini_api.py index c8cab8b..46f4817 100644 --- a/pentestgpt/utils/APIs/gemini_api.py +++ b/pentestgpt/utils/APIs/gemini_api.py @@ -10,8 +10,11 @@ import google.generativeai as genai import loguru import tiktoken -from google.generativeai.types import (HarmBlockThreshold, HarmCategory, - SafetySettingDict) +from google.generativeai.types import ( + HarmBlockThreshold, + HarmCategory, + SafetySettingDict, +) from langfuse.model import InitialGeneration, Usage from tenacity import * diff --git a/pentestgpt/utils/pentest_gpt.py b/pentestgpt/utils/pentest_gpt.py index 34ec545..8acd750 100644 --- a/pentestgpt/utils/pentest_gpt.py +++ b/pentestgpt/utils/pentest_gpt.py @@ -17,9 +17,12 @@ from pentestgpt.utils.APIs.module_import import dynamic_import from pentestgpt.utils.chatgpt import ChatGPT from pentestgpt.utils.prompt_select import prompt_ask, prompt_select -from pentestgpt.utils.task_handler import (local_task_entry, - localTaskCompleter, main_task_entry, - mainTaskCompleter) +from pentestgpt.utils.task_handler import ( + local_task_entry, + localTaskCompleter, + main_task_entry, + mainTaskCompleter, +) from pentestgpt.utils.web_parser import google_search logger = loguru.logger diff --git a/pentestgpt/utils/pentest_gpt_rebuilt.py b/pentestgpt/utils/pentest_gpt_rebuilt.py index 785526a..1459045 100644 --- a/pentestgpt/utils/pentest_gpt_rebuilt.py +++ b/pentestgpt/utils/pentest_gpt_rebuilt.py @@ -17,9 +17,12 @@ from pentestgpt.utils.APIs.module_import import dynamic_import from pentestgpt.utils.chatgpt import ChatGPT from pentestgpt.utils.prompt_select import prompt_ask, prompt_select -from pentestgpt.utils.task_handler import (local_task_entry, - localTaskCompleter, main_task_entry, - mainTaskCompleter) +from pentestgpt.utils.task_handler import ( + local_task_entry, + localTaskCompleter, + main_task_entry, + mainTaskCompleter, +) from pentestgpt.utils.web_parser import google_search logger = loguru.logger diff --git a/pentestgpt/utils/prompt_select.py b/pentestgpt/utils/prompt_select.py index 7fee7d4..94ba3e4 100644 --- a/pentestgpt/utils/prompt_select.py +++ b/pentestgpt/utils/prompt_select.py @@ -3,8 +3,7 @@ from prompt_toolkit.application import Application from prompt_toolkit.formatted_text import HTML from prompt_toolkit.key_binding.defaults import load_key_bindings -from prompt_toolkit.key_binding.key_bindings import (KeyBindings, - merge_key_bindings) +from prompt_toolkit.key_binding.key_bindings import KeyBindings, merge_key_bindings from prompt_toolkit.layout import Layout from prompt_toolkit.layout.containers import HSplit from prompt_toolkit.shortcuts import prompt