From 27dfd7706dd7207e1a62a81456787520a53e8c34 Mon Sep 17 00:00:00 2001 From: Grey_D Date: Thu, 27 Apr 2023 18:06:03 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20minor=20logic=20improvem?= =?UTF-8?q?ent?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add exception handling in main loop so that logs can also be saved even when error occurs. --- utils/pentest_gpt.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/utils/pentest_gpt.py b/utils/pentest_gpt.py index d53c4ee..b46a82a 100644 --- a/utils/pentest_gpt.py +++ b/utils/pentest_gpt.py @@ -481,11 +481,19 @@ class pentestGPT: # 4. enter the main loop. while True: - result = self.input_handler() - self.console.print( - "-----------------------------------------", style="bold white" - ) - if not result: # end the session + try: + result = self.input_handler() + self.console.print( + "-----------------------------------------", style="bold white" + ) + if not result: # end the session + break + except Exception as e: # catch all general exception. + # log the exception + self.log_conversation("exception", "e") + # print the exception + self.console.print("Exception: " + str(e), style="bold red") + # safely quit the session break # Summarize the session and end