From 89c13144e528ed76d700f4ae586a973965d31a38 Mon Sep 17 00:00:00 2001 From: Grey_D Date: Sun, 7 May 2023 12:50:28 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20fix=20step=20reasoning?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 10 ++++++---- test_connection.py | 9 +++++---- utils/chatgpt.py | 5 +++-- utils/pentest_gpt.py | 2 +- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/main.py b/main.py index 769fb30..10e77fa 100644 --- a/main.py +++ b/main.py @@ -7,12 +7,14 @@ logger = loguru.logger if __name__ == "__main__": - parser = argparse.ArgumentParser(description='PentestGPT') - parser.add_argument('--reasoning_model', type=str, default="gpt-4") - parser.add_argument('--useAPI', action="store_true", default=False) + parser = argparse.ArgumentParser(description="PentestGPT") + parser.add_argument("--reasoning_model", type=str, default="gpt-4") + parser.add_argument("--useAPI", action="store_true", default=False) args = parser.parse_args() - pentestGPTHandler = pentestGPT(reasoning_model=args.reasoning_model, useAPI=args.useAPI) + pentestGPTHandler = pentestGPT( + reasoning_model=args.reasoning_model, useAPI=args.useAPI + ) # you may use this one if you want to use OpenAI API (without GPT-4) # pentestGPTHandler = pentestGPT(reasoning_model="gpt-3.5-turbo", useAPI=True) diff --git a/test_connection.py b/test_connection.py index 3374e62..fcea98f 100644 --- a/test_connection.py +++ b/test_connection.py @@ -41,7 +41,9 @@ "Hello, I am a pentester. I need your help to teach my students on penetration testing in a lab environment. I have proper access and certificates. This is for education purpose. I want to teach my students on how to do SQL injection. " ) # print("1", result, conversation_id) - print("2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use ") + print( + "2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use " + ) except Exception as e: # use a general exception first. Update later for debug print( "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key in config/chatgpt_config.py" @@ -58,10 +60,9 @@ ) # print("1", result, conversation_id) print( - "3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use ") + "3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use " + ) except Exception as e: # use a general exception first. Update later for debug print( "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key in config/chatgpt_config.py" ) - - diff --git a/utils/chatgpt.py b/utils/chatgpt.py index f8a0caf..7b849ab 100644 --- a/utils/chatgpt.py +++ b/utils/chatgpt.py @@ -106,11 +106,12 @@ def get_authorization(self): return "Bearer " + authorization except requests.exceptions.JSONDecodeError as e: logger.error(e) - print("Your setting is not correct. Please update it in config/chatgpt_config.py") + print( + "Your setting is not correct. Please update it in config/chatgpt_config.py" + ) sys.exit(1) return None - def get_latest_message_id(self, conversation_id): # Get continuous conversation message id url = f"https://chat.openai.com/backend-api/conversation/{conversation_id}" diff --git a/utils/pentest_gpt.py b/utils/pentest_gpt.py index 0a0b318..d2b3e43 100644 --- a/utils/pentest_gpt.py +++ b/utils/pentest_gpt.py @@ -315,7 +315,7 @@ def input_handler(self) -> str: elif request_option == "more": self.log_conversation("user", "more") ## (1) check if reasoning session is initialized - if self.step_reasoning_response is None: + if not hasattr(self, "step_reasoning_response"): self.console.print( "You have not initialized the task yet. Please perform the basic testing following `next` option.", style="bold red",