docs: ✏️ update documentation for pentestgpt design

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