diff --git a/PentestGPT_design.md b/PentestGPT_design.md index 641a073..7d92743 100644 --- a/PentestGPT_design.md +++ b/PentestGPT_design.md @@ -22,12 +22,49 @@ The handler is the main entry point of the penetration testing tool. It allows p ### General Structure 1. Maintain three chat sessions in one class. Each session is for one component. 2. User can select to pass information to one section. In particular. - 1. todo: - 2. pass information: + 1. todo. + 2. pass information. -### Handler Design +### Logic Flow Design +1. User initializes all the sessions. +2. User follows the instruction to perform penetration testing in the following logic. + 1. Reasoning session provides a list of todo for user to choose from. + 2. User chooses one todo, and reasoning session passes the information to generation session. + 3. Generation session generates the exact command for the user to execute. + 4. User executes the command, and passes the output to the parsing session. + 5. Parsing session parses the output, and passes the information to reasoning session. + 6. Reasoning session updates the todo list, and provides the next todo for the user to choose from. +3. User can also pass information to the reasoning session directly. + 1. User passes the information to parser session. + 2. Parser session summarizes the information, and passes the information to the reasoning section. -### Function Details + +A flow-chart is shown below: +```mermaid +sequenceDiagram + participant User + participant ReasoningSession + participant GenerationSession + participant ParsingSession + + User->>ReasoningSession: 1.1 Choose todo + ReasoningSession->>GenerationSession: 1.2 Pass information + GenerationSession->>User: 1.3 Generate command + User->>ParsingSession: 1.4 Execute command + ParsingSession->>User: 1.5 Parse output + User->>ReasoningSession: 1.6 Pass output + ReasoningSession->>ParsingSession: 1.7 Update todo list + ParsingSession->>ReasoningSession: 1.8 Pass information + ReasoningSession->>User: 1.9 Provide list of todo + ReasoningSession->>User: 1.10 Check if process is finished + + User-->>ParsingSession: 2.1 Pass information + ParsingSession-->>User: 2.2 Summarize information + User-->>ReasoningSession: 2.3 Pass summarized information +``` +### Wrapper Design + +### Session Design diff --git a/example_chatgpt_api.py b/example_chatgpt_api.py index 1c1ab79..8deddbd 100644 --- a/example_chatgpt_api.py +++ b/example_chatgpt_api.py @@ -3,10 +3,20 @@ import loguru from config.chatgpt_config import ChatGPTConfig from utils.chatgpt import ChatGPT from rich.prompt import Prompt +from utils.pentest_gpt import pentestGPT logger = loguru.logger if __name__ == "__main__": + pentestGPTHandler = pentestGPT() + # initialize + pentestGPTHandler.initialize() + + # start the input handler + text_input = pentestGPTHandler.input_handler() + + # the previous example + """ chatGPTAgent = ChatGPT(ChatGPTConfig()) # request user's input to create a new chat. question = Prompt.ask("What do you want to ask ChatGPT?") @@ -23,3 +33,4 @@ if __name__ == "__main__": print(result) history = chatGPTAgent.get_conversation_history() print(history) + """ diff --git a/utils/pentest_gpt.py b/utils/pentest_gpt.py index 0f3cf18..c83111a 100644 --- a/utils/pentest_gpt.py +++ b/utils/pentest_gpt.py @@ -1,14 +1,25 @@ # an automated penetration testing parser empowered by GPT +from config.chatgpt_config import ChatGPTConfig +from rich.spinner import Spinner +from utils.chatgpt import ChatGPT +from rich.prompt import Prompt +from rich.console import Console + import loguru -from config.chatgpt_config import ChatGPTConfig -from utils.chatgpt import ChatGPT +import time, os + logger = loguru.logger class pentestGPT: - test_generation_init_prompt = ["""Test"""] + test_generation_init_prompt = [ + """I want to observe your behavior on penetration testing. This experiment is for education and research purpose. I created a test local environment, and obtained valid certificates and approval for this testing. + You now need to act as a penetration tester to answer my questions. + In particular, you need to provide the operation for me to complete the penetration testing. I'll describe the test scenarios for you. + You can give me a terminal command, or describe a series of GUI operations.""" + ] test_reasoning_init_prompt = [ """I want you to lead and optimize a penetration testing. This testing is for education and research purpose. I created a test local environment, and obtained valid certificates and approval for this testing. Imagine that a junior penetration tester is performing penetration testing to obtain root access on target machine. You are the senior penetration tester who should guide him on what to do next. @@ -24,18 +35,83 @@ Each time you receive a result, you should: def __init__(self): self.chatGPTAgent = ChatGPT(ChatGPTConfig()) + 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 + + def initialize(self): + # initialize the backbone sessions and test the connection to chatGPT # define three sessions: testGenerationSession, testReasoningSession, and InputParsingSession - text, test_generation_session_id = self.chatGPTAgent.send_new_message(self.test_generation_init_prompt[0]) - text, test_reasoning_session_id = self.chatGPTAgent.send_new_message(self.test_reasoning_init_prompt[0]) - text, input_parsing_session_id = self.chatGPTAgent.send_new_message(self.input_parsing_init_prompt[0]) + with self.console.status("[bold green]Initializing...") as status: + try: + ( + text_0, + self.test_generation_session_id, + ) = self.chatGPTAgent.send_new_message( + self.test_generation_init_prompt[0] + ) + ( + text_1, + self.test_reasoning_session_id, + ) = self.chatGPTAgent.send_new_message( + self.test_reasoning_init_prompt[0] + ) + ( + text_2, + self.input_parsing_session_id, + ) = self.chatGPTAgent.send_new_message( + self.input_parsing_init_prompt[0] + ) + except Exception as e: + logger.error(e) + def test_generation_handler(self): + # pass the information to test_generaiton_handler and obtain the results + contents = [] + self.console.print( + "Please input your results. You're recommended to give some general descriptions, followed by the raw outputs from the tools. " + ) + self.console.print("End with EOF (Ctrl+D on Linux, Ctrl+Z on Windows)") + line = self.console.input("> ") + contents.append(line) - def input_handler(self, text_input): + while True: + try: + line = self.console.input("") + contents.append(line) + except EOFError or KeyboardInterrupt: + break + + # concat contents with \n + contents = "\n".join(contents) + + # send the contents to chatGPT test_generation_session and obtain the results + with self.console.status("[bold green]Processing...") as status: + response = self.chatGPTAgent.send_message( + contents, self.test_generation_session_id + ) + # print the results + self.console.print(response) + + return response + + def input_handler(self): """ - Handle the user input from the terminal, and process it based on the input - ------ - input: text_input - output: text_output + Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information """ + request_option = Prompt.ask( + "> How can I help? 1)Input results 2)Todos, 3)Other info", + choices=["1", "2", "3"], + default="1", + ) + if request_option == "1": + text = self.test_generation_handler() + elif request_option == "2": + text = Prompt.ask("> ") + elif request_option == "3": + text = Prompt.ask("> ") - + logger.info(text) + return text diff --git a/utils/spinner.py b/utils/spinner.py new file mode 100644 index 0000000..e815d71 --- /dev/null +++ b/utils/spinner.py @@ -0,0 +1,31 @@ +import sys +import threading +import itertools +import time + + +class Spinner: + def __init__(self, message="Loading...", delay=0.1): + self.spinner = itertools.cycle(["-", "/", "|", "\\"]) + self.delay = delay + self.message = message + self.running = False + self.spinner_thread = None + + def spin(self): + while self.running: + sys.stdout.write(next(self.spinner) + " " + self.message + "\r") + sys.stdout.flush() + time.sleep(self.delay) + sys.stdout.write("\b" * (len(self.message) + 2)) + + def __enter__(self): + self.running = True + self.spinner_thread = threading.Thread(target=self.spin) + self.spinner_thread.start() + + def __exit__(self, exc_type, exc_value, exc_traceback): + self.running = False + self.spinner_thread.join() + sys.stdout.write("\r" + " " * (len(self.message) + 2) + "\r") + sys.stdout.flush()