# an automated penetration testing parser empowered by GPT import json import os import sys import textwrap import time import traceback import loguru from prompt_toolkit.formatted_text import HTML from prompt_toolkit.shortcuts import confirm from rich.console import Console from rich.spinner import Spinner from pentestgpt.config.chat_config import ChatGPTConfig from pentestgpt.prompts.prompt_class import PentestGPTPrompt 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.web_parser import google_search logger = loguru.logger def prompt_continuation(width, line_number, wrap_count): """ The continuation: display line numbers and '->' before soft wraps. Notice that we can return any kind of formatted text from here. The prompt continuation doesn't have to be the same width as the prompt which is displayed before the first line, but in this example we choose to align them. The `width` input that we receive here represents the width of the prompt. """ if wrap_count > 0: return " " * (width - 3) + "-> " text = ("- %i - " % (line_number + 1)).rjust(width) return HTML("%s") % text class pentestGPT: postfix_options = { "tool": "The input content is from a security testing tool. You need to list down all the points that are interesting to you; you should summarize it as if you are reporting to a senior penetration tester for further guidance.\n", "user-comments": "The input content is from user comments.\n", "web": "The input content is from web pages. You need to summarize the readable-contents, and list down all the points that can be interesting for penetration testing.\n", "default": "The user did not specify the input source. You need to summarize based on the contents.\n", } options_desc = { "tool": " Paste the output of the security test tool used", "user-comments": "", "web": " Paste the relevant content of a web page", "default": " Write whatever you want, the tool will handle it", } def __init__( self, log_dir="logs", reasoning_model="gpt-4-turbo", parsing_model="gpt-4-turbo", useAPI=True, azure=False, use_langfuse_logging=False, ): self.log_dir = log_dir logger.add(sink=os.path.join(log_dir, "pentestGPT.log")) self.save_dir = "test_history" self.task_log = ( {} ) # the information that can be saved to continue in the next session self.useAPI = useAPI self.parsing_char_window = 16000 # the chunk size for parsing in # of chars # TODO: link the parsing_char_window to the model used # load the module reasoning_model_object = dynamic_import( reasoning_model, self.log_dir, use_langfuse_logging=use_langfuse_logging ) generation_model_object = dynamic_import( reasoning_model, self.log_dir, use_langfuse_logging=use_langfuse_logging ) parsing_model_object = dynamic_import( parsing_model, self.log_dir, use_langfuse_logging=use_langfuse_logging ) if useAPI is False: # deprecated usage of cookie self.parsingAgent = ChatGPT(ChatGPTConfig(log_dir=self.log_dir)) self.reasoningAgent = ChatGPT( ChatGPTConfig(model=reasoning_model, log_dir=self.log_dir) ) else: self.parsingAgent = parsing_model_object self.generationAgent = generation_model_object self.reasoningAgent = reasoning_model_object self.prompts = PentestGPTPrompt self.console = Console() self.spinner = Spinner("line", "Processing") self.test_generation_session_id = None self.test_reasoning_session_id = None self.input_parsing_session_id = None self.chat_count = 0 self.step_reasoning = ( None # the response from the reasoning session for the current step ) self.history = { "user": [], "pentestGPT": [], "reasoning": [], "input_parsing": [], "generation": [], "exception": [], } # the history of the current conversation # print the initialization message on the current implementation. self.console.print( "Welcome to pentestGPT, an automated penetration testing parser empowered by GPT.", style="bold green", ) self.console.print("The settings are: ") 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") def log_conversation(self, source, text): """ append the conversation into the history Parameters: ---------- source: str the source of the conversation text: str the content of the conversation """ # append the conversation into the history timestamp = time.time() if source not in self.history.keys(): # an exception source = "exception" self.history[source].append((timestamp, text)) def refresh_session(self): if self.useAPI: self.console.print( "You're using API mode, so no need to refresh the session." ) self.log_conversation( "pentestGPT", "You're using API mode, so no need to refresh the session.", ) else: self.console.print( "Please ensure that you put the curl command into `config/chatgpt_config_curl.txt`", style="bold green", ) self.log_conversation( "pentestGPT", "Please ensure that you put the curl command into `config/chatgpt_config_curl.txt`", ) input("Press Enter to continue...") self.parsingAgent.refresh() self.reasoningAgent.refresh() self.console.print( "Session refreshed. If you receive the same session refresh request, please refresh the ChatGPT page and paste the new curl request again.", style="bold green", ) self.log_conversation("pentestGPT", "Session refreshed.") return "Session refreshed." def _feed_init_prompts(self): # 1. User firstly provide basic information of the task init_description = prompt_ask( "Please describe the penetration testing task in one line, including the target IP, task type, etc.\n> ", multiline=False, ) self.log_conversation("user", init_description) self.task_log["task description"] = init_description # 2. Provide the information to the reasoning session for the task initialization. # Note that this information is not parsed by the three-step process in reasoning. # It is directly used to initialize the task. prefixed_init_description = self.prompts.task_description + init_description with self.console.status( "[bold green] Constructing Initial Penetration Testing Tree..." ) as status: _reasoning_response = self.reasoningAgent.send_message( prefixed_init_description, self.test_reasoning_session_id ) # 3. Pass to generation session for more details. # Note that the generation session is not used for the task initialization. with self.console.status("[bold green] Generating Initial Task") as status: _generation_response = self.generationAgent.send_message( self.prompts.todo_to_command + _reasoning_response, self.test_generation_session_id, ) # Display the initial generation result response = _reasoning_response + "\n" + _generation_response self.console.print("PentestGPT output: ", style="bold green") self.console.print(response) self.log_conversation("PentestGPT", "PentestGPT output:" + response) def initialize(self, previous_session_ids=None): # initialize the backbone sessions and test the connection to chatGPT # define three sessions: testGenerationSession, testReasoningSession, and InputParsingSession if previous_session_ids is not None and self.useAPI is False: self.test_generation_session_id = previous_session_ids.get( "test_generation", None ) self.test_reasoning_session_id = previous_session_ids.get("reasoning", None) self.input_parsing_session_id = previous_session_ids.get("parsing", None) # debug the three sessions print(f"Previous session ids: {str(previous_session_ids)}") print(f"Test generation session id: {str(self.test_generation_session_id)}") print(f"Test reasoning session id: {str(self.test_reasoning_session_id)}") print(f"Input parsing session id: {str(self.input_parsing_session_id)}") print("-----------------") self.task_log = previous_session_ids.get("task_log", {}) self.console.print(f"Task log: {str(self.task_log)}", style="bold green") print("You may use discussion function to remind yourself of the task.") ## verify that all the sessions are not None if ( self.test_generation_session_id is None or self.test_reasoning_session_id is None or self.input_parsing_session_id is None ): self.console.print( "[bold red] Error: the previous session ids are not valid. Loading new sessions" ) self.initialize() else: with self.console.status( "[bold green] Initialize ChatGPT Sessions..." ) as status: try: ( text_0, self.test_generation_session_id, ) = self.generationAgent.send_new_message( self.prompts.generation_session_init, ) ( text_1, self.test_reasoning_session_id, ) = self.reasoningAgent.send_new_message( self.prompts.reasoning_session_init ) ( text_2, self.input_parsing_session_id, ) = self.parsingAgent.send_new_message( self.prompts.input_parsing_init ) except Exception as e: logger.error(e) self.console.print("- ChatGPT Sessions Initialized.", style="bold green") self._feed_init_prompts() def reasoning_handler(self, text) -> str: # summarize the contents if necessary. if len(text) > self.parsing_char_window: text = self.input_parsing_handler(text) """ # pass the information to reasoning_handler and obtain the results response = self.reasoningAgent.send_message( self.prompts.process_results + text, self.test_reasoning_session_id ) # log the conversation """ # PentestGPT Reasoning Logic ## 1. Given the information, update the PTT _updated_ptt_response = self.reasoningAgent.send_message( self.prompts.process_results + text, self.test_reasoning_session_id ) ## 2. Validate if the PTT is correct # TODO ## 3. If the PTT is correct, select all the to-dos _task_selection_response = self.reasoningAgent.send_message( self.prompts.process_results_task_selection, self.test_reasoning_session_id ) # get the complete output: response = _updated_ptt_response + _task_selection_response self.log_conversation("reasoning", response) return response def input_parsing_handler(self, text, source=None) -> str: prefix = "Please summarize the following input. " # do some engineering trick here. Add postfix to the input to make it more understandable by LLMs. if source is not None and source in self.postfix_options.keys(): prefix += self.postfix_options[source] # The default token-size limit is 4096 (web UI even shorter). 1 token ~= 4 chars in English # Use textwrap to split inputs. Limit to 2000 token (8000 chars) for each input # (1) replace all the newlines with spaces text = text.replace("\r", " ").replace("\n", " ") # (2) wrap the text wrapped_text = textwrap.fill(text, 8000) wrapped_inputs = wrapped_text.split("\n") # (3) send the inputs to chatGPT input_parsing_session and obtain the results summarized_content = "" for wrapped_input in wrapped_inputs: word_limit = f"Please ensure that the input is less than {8000 / len(wrapped_inputs)} words.\n" summarized_content += self.parsingAgent.send_message( prefix + word_limit + wrapped_input, self.input_parsing_session_id ) # log the conversation self.log_conversation("input_parsing", summarized_content) return summarized_content def test_generation_handler(self, text): # send the contents to chatGPT test_generation_session and obtain the results response = self.generationAgent.send_message( text, self.test_generation_session_id ) # log the conversation self.log_conversation("generation", response) return response def local_input_handler(self) -> str: """ Request for user's input to handle the local task """ local_task_response = "" self.chat_count += 1 local_request_option = local_task_entry() self.log_conversation("user", local_request_option) if local_request_option == "help": print(localTaskCompleter().task_details) elif local_request_option == "discuss": ## (1) Request for user multi-line input self.console.print( "Please share your findings and questions with PentestGPT." ) self.log_conversation( "pentestGPT", "Please share your findings and questions with PentestGPT. (End with )", ) user_input = prompt_ask("Your input: ", multiline=True) self.log_conversation("user", user_input) ## (2) pass the information to the reasoning session. with self.console.status("[bold green] PentestGPT Thinking...") as status: local_task_response = self.test_generation_handler( self.prompts.local_task_prefix + user_input ) ## (3) print the results self.console.print("PentestGPT:\n", style="bold green") self.console.print(local_task_response + "\n", style="yellow") self.log_conversation("pentestGPT", local_task_response) elif local_request_option == "brainstorm": ## (1) Request for user multi-line input self.console.print( "Please share your concerns and questions with PentestGPT." ) self.log_conversation( "pentestGPT", "Please share your concerns and questions with PentestGPT. End with )", ) user_input = prompt_ask("Your input: ", multiline=True) self.log_conversation("user", user_input) ## (2) pass the information to the reasoning session. with self.console.status("[bold green] PentestGPT Thinking...") as status: local_task_response = self.test_generation_handler( self.prompts.local_task_brainstorm + user_input ) ## (3) print the results self.console.print("PentestGPT:\n", style="bold green") self.console.print(local_task_response + "\n", style="yellow") self.log_conversation("pentestGPT", local_task_response) elif local_request_option == "google": # get the users input self.console.print( "Please enter your search query. PentestGPT will summarize the info from google. (End with ) ", style="bold green", ) self.log_conversation( "pentestGPT", "Please enter your search query. PentestGPT will summarize the info from google.", ) user_input = prompt_ask("Your input: ", multiline=False) self.log_conversation("user", user_input) with self.console.status("[bold green] PentestGPT Thinking...") as status: # query the question result: dict = google_search(user_input, 5) # 5 results by default # summarize the results # TODO local_task_response = ( "Google search results:\n" + "still under development." ) self.console.print(local_task_response + "\n", style="yellow") self.log_conversation("pentestGPT", local_task_response) return local_task_response elif local_request_option == "continue": self.console.print("Exit the local task and continue the main task.") self.log_conversation( "pentestGPT", "Exit the local task and continue the main task." ) local_task_response = "continue" return local_task_response def input_handler(self) -> str: """ Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information (discuss), (4) google. (4) end. The design details are based on PentestGPT_design.md Return ----- response: str The response from the chatGPT model. """ self.chat_count += 1 request_option = main_task_entry() self.log_conversation("user", request_option) # always check if the session expires. # check if session expires if not self.useAPI: conversation_history = self.parsingAgent.get_conversation_history() while conversation_history is None: self.refresh_session() conversation_history = self.parsingAgent.get_conversation_history() if request_option == "help": print(mainTaskCompleter().task_details) if request_option == "next": ## (1) pass the information to input_parsing session. ## Give an option list for user to choose from options = list(self.postfix_options.keys()) opt_desc = list(self.options_desc.values()) value_list = [ ( i, HTML( f'' ), ) for i in range(len(options)) ] source = prompt_select( title="Please choose the source of the information.", values=value_list ) self.console.print( "Your input: (End with )", style="bold green" ) user_input = prompt_ask("> ", multiline=True) self.log_conversation( "user", f"Source: {options[int(source)]}" + "\n" + user_input ) with self.console.status("[bold green] PentestGPT Thinking...") as status: parsed_input = self.input_parsing_handler( user_input, source=options[int(source)] ) ## (2) pass the summarized information to the reasoning session. reasoning_response = self.reasoning_handler(parsed_input) self.step_reasoning_response = reasoning_response ## (3) print the results self.console.print( "Based on the analysis, the following tasks are recommended:", style="bold green", ) self.console.print(reasoning_response + "\n") self.log_conversation( "pentestGPT", "Based on the analysis, the following tasks are recommended:" + reasoning_response, ) response = reasoning_response elif request_option == "more": self.log_conversation("user", "more") ## (1) check if reasoning session is initialized 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", ) response = "You have not initialized the task yet. Please perform the basic testing following `next` option." self.log_conversation("pentestGPT", response) return response ## (2) start local task generation. ### (2.1) ask the reasoning session to analyze the current situation, and explain the task self.console.print( "PentestGPT will generate more test details, and enter the sub-task generation mode. (Pressing Enter to continue)", style="bold green", ) self.log_conversation( "pentestGPT", "PentestGPT will generate more test details, and enter the sub-task generation mode.", ) input() ### (2.2) pass the sub-tasks to the test generation session with self.console.status("[bold green] PentestGPT Thinking...") as status: generation_response = self.test_generation_handler( self.step_reasoning_response ) _local_init_response = self.test_generation_handler( self.prompts.local_task_init ) self.console.print( "Below are the further details.", style="bold green", ) self.console.print(generation_response + "\n") response = generation_response self.log_conversation("pentestGPT", response) ### (2.3) local task handler while True: local_task_response = self.local_input_handler() if local_task_response == "continue": # break the local task handler break elif request_option == "todo": ## log that user is asking for todo list self.log_conversation("user", "todo") ## (1) ask the reasoning session to analyze the current situation, and list the top sub-tasks with self.console.status("[bold green] PentestGPT Thinking...") as status: reasoning_response = self.reasoning_handler(self.prompts.ask_todo) ## (2) pass the sub-tasks to the test_generation session. message = self.prompts.todo_to_command + "\n" + reasoning_response generation_response = self.test_generation_handler(message) ## (3) print the results self.console.print( "Based on the analysis, the following tasks are recommended:", style="bold green", ) self.console.print(reasoning_response + "\n") self.console.print( "You can follow the instructions below to complete the tasks.", style="bold green", ) self.console.print(generation_response + "\n") response = reasoning_response self.log_conversation( "pentestGPT", ( ( ( ( "Based on the analysis, the following tasks are recommended:" + response ) + "\n" ) + "You can follow the instructions below to complete the tasks." ) + generation_response ), ) elif request_option == "discuss": ## (1) Request for user multi-line input self.console.print( "Please share your thoughts/questions with PentestGPT. (End with ) " ) self.log_conversation( "pentestGPT", "Please share your thoughts/questions with PentestGPT." ) user_input = prompt_ask("Your input: ", multiline=True) self.log_conversation("user", user_input) ## (2) pass the information to the reasoning session. with self.console.status("[bold green] PentestGPT Thinking...") as status: response = self.reasoning_handler(self.prompts.discussion + user_input) ## (3) print the results self.console.print("PentestGPT:\n", style="bold green") self.console.print(response + "\n", style="yellow") self.log_conversation("pentestGPT", response) elif request_option == "google": # get the users input self.console.print( "Please enter your search query. PentestGPT will summarize the info from google. (End with ) ", style="bold green", ) self.log_conversation( "pentestGPT", "Please enter your search query. PentestGPT will summarize the info from google.", ) user_input = prompt_ask("Your input: ", multiline=False) self.log_conversation("user", user_input) with self.console.status("[bold green] PentestGPT Thinking...") as status: # query the question result: dict = google_search(user_input, 5) # 5 results by default # summarize the results # TODO response = "Google search results:\n" + "still under development." self.console.print(response + "\n", style="yellow") self.log_conversation("pentestGPT", response) return response elif request_option == "quit": response = False self.console.print("Thank you for using PentestGPT!", style="bold green") self.log_conversation("pentestGPT", "Thank you for using PentestGPT!") else: self.console.print("Please key in the correct options.", style="bold red") self.log_conversation("pentestGPT", "Please key in the correct options.") response = "Please key in the correct options." return response def save_session(self): """ Save the current session for next round of usage. The test information is saved in the directory `./test_history` """ self.console.print( "Before you quit, you may want to save the current session.", style="bold green", ) # 1. Require a save name from the user. If not, use the current time as the save name. save_name = prompt_ask( "Please enter the name of the current session. (Default with current timestamp)\n> ", multiline=False, ) if save_name == "": save_name = str(time.time()) # 2. Save the current session with open( os.path.join( os.path.realpath(os.path.dirname(__file__)), os.pardir, os.pardir, self.save_dir, save_name, ), "w", ) as f: # store the three ids and task_log session_ids = { "reasoning": self.test_reasoning_session_id, "test_generation": self.test_generation_session_id, "parsing": self.input_parsing_session_id, "task_log": self.task_log, } json.dump(session_ids, f) self.console.print( f"The current session is saved as {save_name}", style="bold green" ) return def _preload_session(self) -> dict: """ Preload the session from the save directory. Returns: dict: the session ids for the three sessions. None if no previous session is found. """ if continue_from_previous := confirm( "Do you want to continue from previous session?" ): # load the filenames from the save directory filenames = os.listdir( os.path.join( os.path.realpath(os.path.dirname(__file__)), os.pardir, os.pardir, self.save_dir, ) ) if len(filenames) == 0: print("No previous session found. Please start a new session.") return None else: # print all the files print("Please select the previous session by its index (integer):") for i, filename in enumerate(filenames): print(f"{str(i)}. {filename}") # ask for the user input try: previous_testing_name = filenames[ int(input("Please key in your option (integer): ")) ] print(f"You selected: {previous_testing_name}") except ValueError as e: print("You input an invalid option. Will start a new session.") return None elif continue_from_previous is False: return None else: print("You input an invalid option. Will start a new session.") return None # 2. load the previous session information if previous_testing_name is not None: # try to load the file content with json try: with open( os.path.join( os.path.realpath(os.path.dirname(__file__)), os.pardir, os.pardir, self.save_dir, previous_testing_name, ), "r", ) as f: return json.load(f) except Exception as e: print( "Error when loading the previous session. The file name is not correct" ) print(e) previous_testing_name = None return None def main(self): """ The main function of pentestGPT. The design is based on PentestGPT_design.md """ # 0. initialize the backbone sessions and test the connection to chatGPT loaded_ids = self._preload_session() self.initialize(previous_session_ids=loaded_ids) # enter the main loop. while True: 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", str(e)) # print the exception self.console.print(f"Exception: {str(e)}", style="bold red") # add a more detailed debugging exc_type, exc_obj, exc_tb = sys.exc_info() fname = os.path.split(exc_tb.tb_frame.f_code.co_filename)[1] self.console.print( "Exception details are below. You may submit an issue on github and paste the error trace", style="bold green", ) # self.console.print(exc_type, fname, exc_tb.tb_lineno) print(traceback.format_exc()) # safely quit the session break # log the session. Save self.history into a txt file based on timestamp timestamp = time.time() log_name = f"pentestGPT_log_{str(timestamp)}.txt" # save it in the logs folder log_path = os.path.join(self.log_dir, log_name) with open(log_path, "w") as f: json.dump(self.history, f) # save the sessions; continue from previous testing self.save_session() if __name__ == "__main__": pentestGPT = pentestGPT() pentestGPT.main()