docs: ✏️ update documentation for pentestgpt design

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
 
 
 

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)
+    """

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

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()