diff --git a/pentestgpt/test_connection.py b/pentestgpt/test_connection.py index 1dcc905..0bf9ebd 100644 --- a/pentestgpt/test_connection.py +++ b/pentestgpt/test_connection.py @@ -33,28 +33,7 @@ def main(): chatgpt_config = ChatGPTConfig(api_base=args.baseUrl) console = Console() - # 1. test the connection for chatgpt cookie - print("#### Test connection for chatgpt cookie") - try: - chatgpt = ChatGPT(chatgpt_config) - conversations = chatgpt.get_conversation_history() - if conversations is not None: - console.print( - "1. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use ", - style="bold green", - ) - else: - console.print( - "The cookie is not properly configured with ChatGPT Cookie. Please follow README to update cookie through `export CHATGPT_COOKIE=`", - style="bold red", - ) - except Exception as e: # use a general exception first. Update later for debug - logger.error(e) - print( - "The cookie is not properly configured. Please follow README to update cookie in config/chatgpt_config.py" - ) - - # 2. test the connection for chatgpt api with GPT-4 + # 1. test the connection for chatgpt api with GPT-4 print("#### Test connection for OpenAI api (GPT-4)") try: chatgpt_config.model = "gpt-4" @@ -62,35 +41,54 @@ def main(): openai.api_key = chatgpt_config.openai_key result, conversation_id = chatgpt.send_new_message("Hi how are you?") console.print( - "2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use ", + "1. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use ", style="bold green", ) except Exception as e: # use a general exception first. Update later for debug console.print( - "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key through `export OPENAI_KEY=<>`.", + "1. The OpenAI API key is not properly configured. Please follow README to update OpenAI API key through `export OPENAI_KEY=<>`.", style="bold red", ) print("The error is below:", e) - # 3. test the connection for chatgpt api with GPT-3.5 + # 2. 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" + 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( - "3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use ", + "2. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use ", style="bold green", ) except Exception as e: # use a general exception first. Update later for debug logger.error(e) console.print( - "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key through `export OPENAI_KEY=<>`", + "2. 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_KEY=<>`", style="bold red", ) print("The error is below:", e) + # 3. test the connection for chatgpt cookie (deprecated) + print("#### Test connection for chatgpt cookie") + try: + chatgpt = ChatGPT(chatgpt_config) + conversations = chatgpt.get_conversation_history() + if conversations is not None: + console.print( + "3. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use ", + style="bold green", + ) + else: + console.print( + "3. The cookie is not properly configured with ChatGPT Cookie. If you're not using cookie bypass for testing, please neglect this message.", + style="bold red", + ) + except Exception as e: # use a general exception first. Update later for debug + logger.error(e) + print("The cookie is not properly configured.") + if __name__ == "__main__": main() diff --git a/pentestgpt/utils/APIs/chatgpt_api.py b/pentestgpt/utils/APIs/chatgpt_api.py index 0dabac4..e6a0699 100644 --- a/pentestgpt/utils/APIs/chatgpt_api.py +++ b/pentestgpt/utils/APIs/chatgpt_api.py @@ -43,6 +43,7 @@ def __eq__(self, other): class ChatGPTAPI(LLMAPI): def __init__(self, config_class): + self.name = str(config_class.model) openai.api_key = os.getenv("OPENAI_KEY", None) openai.api_base = config_class.api_base self.model = config_class.model diff --git a/pentestgpt/utils/APIs/gpt4all_api.py b/pentestgpt/utils/APIs/gpt4all_api.py index 20e7936..1950c1f 100644 --- a/pentestgpt/utils/APIs/gpt4all_api.py +++ b/pentestgpt/utils/APIs/gpt4all_api.py @@ -45,6 +45,7 @@ def __eq__(self, other): class GPT4ALLAPI(LLMAPI): def __init__(self, config_class): + self.name = str(config_class.model) self.history_length = ( 2 # maintain 2 messages in the history due to gpt4all limitation. ) @@ -71,7 +72,7 @@ def _chat_completion(self, history: List) -> str: if __name__ == "__main__": chatgpt_config = ChatGPTConfig() - chatgpt = ChatGPTAPI() + chatgpt = GPT4ALLAPI() # test is below # 1. create a new conversation diff --git a/pentestgpt/utils/APIs/module_import.py b/pentestgpt/utils/APIs/module_import.py index df06edc..4055f86 100644 --- a/pentestgpt/utils/APIs/module_import.py +++ b/pentestgpt/utils/APIs/module_import.py @@ -70,8 +70,14 @@ def dynamic_import(module_name, log_dir) -> object: return LLM_class_initialized else: - print("Module not found: " + module_name) - return None + print( + "Module not found: " + + module_name + + ". Falling back to use the default gpt-3.5-turbo-16k" + ) + # fall back to gpt-3.5-turbo-16k + LLM_class_initialized = dynamic_import("gpt-3.5-turbo-16k", log_dir) + return LLM_class_initialized if __name__ == "__main__": diff --git a/pentestgpt/utils/llm_api.py b/pentestgpt/utils/llm_api.py index 4f06b15..d3b7e5b 100644 --- a/pentestgpt/utils/llm_api.py +++ b/pentestgpt/utils/llm_api.py @@ -45,6 +45,7 @@ def __eq__(self, other): class LLMAPI: def __init__(self, config: ChatGPTConfig): + self.name = "LLMAPI_base_class" self.config = config openai.api_key = config.openai_key openai.proxy = config.proxies diff --git a/pentestgpt/utils/pentest_gpt.py b/pentestgpt/utils/pentest_gpt.py index d59a16b..2910d7a 100644 --- a/pentestgpt/utils/pentest_gpt.py +++ b/pentestgpt/utils/pentest_gpt.py @@ -95,8 +95,12 @@ def __init__( style="bold green", ) self.console.print("The settings are: ") - self.console.print(f" - parsing model: {parsing_model}", style="bold green") - self.console.print(f" - reasoning model: {reasoning_model}", style="bold green") + self.console.print( + f" - parsing model: {parsing_model_object.name}", style="bold green" + ) + self.console.print( + f" - reasoning model: {reasoning_model_object.name}", style="bold green" + ) self.console.print(f" - use API: {useAPI}", style="bold green") self.console.print(f" - log directory: {log_dir}", style="bold green") diff --git a/setup.py b/setup.py index 2a16f3c..e6356f1 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ setup( name="pentestgpt", - version="0.9.0", + version="0.9.1", description="PentestGPT, a GPT-empowered penetration testing tool", long_description=""" PentestGPT is a penetration testing tool empowered by ChatGPT.