feat: 🎸 v0.3.0. Feature optimization

utils/pentest_gpt.py

@@ -6,6 +6,7 @@ from rich.console import Console
 from prompts.prompt_class import PentestGPTPrompt
 from utils.prompt_select import prompt_select, prompt_ask
 from prompt_toolkit.formatted_text import HTML
+from utils.task_handler import main_task_entry, mainTaskCompleter
 
 import loguru
 import time, os, textwrap
@@ -47,6 +48,10 @@ class pentestGPT:
         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
+        )
 
     def initialize(self):
         # initialize the backbone sessions and test the connection to chatGPT
@@ -122,17 +127,23 @@ class pentestGPT:
         response: str
             The response from the chatGPT model.
         """
-        request_option = prompt_select(
-            title="> Please select your options with cursor: ",
-            values=[
-                ("1", HTML('<style fg="cyan">Input test results</style>')),
-                ("2", HTML('<style fg="cyan">Ask for todos</style>')),
-                ("3", HTML('<style fg="cyan">Discuss with PentestGPT</style>')),
-                ("4", HTML('<style fg="cyan">Exit</style>')),
-            ],
-        )
+        self.chat_count += 1
+
+        request_option = main_task_entry()
+        # request_option = prompt_select(
+        #     title=f"({self.chat_count}) > Please select your options with cursor: ",
+        #     values=[
+        #         ("1", HTML('<style fg="cyan">Input test results</style>')),
+        #         ("2", HTML('<style fg="cyan">Ask for todos</style>')),
+        #         ("3", HTML('<style fg="cyan">Discuss with PentestGPT</style>')),
+        #         ("4", HTML('<style fg="cyan">Exit</style>')),
+        #     ],
+        # )
         # pass output
-        if request_option == "1":
+        if request_option == "help":
+            print(mainTaskCompleter().task_details)
+
+        if request_option == "next":
             ## (1) pass the information to input_parsing session.
             ## Give a option list for user to choose from
             options = list(self.postfix_options.keys())
@@ -153,23 +164,39 @@ class pentestGPT:
                 )
                 ## (2) pass the summarized information to the reasoning session.
                 reasoning_response = self.reasoning_handler(parsed_input)
-                ## (3) pass the reasoning results to the test_generation session.
-                generation_response = self.test_generation_handler(reasoning_response)
-            ## (4) print the results
+                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")
+            response = reasoning_response
+
+        # generate more test details (beginner mode)
+        elif request_option == "more":
+            ## (1) pass the reasoning results to the test_generation session.
+            if self.step_reasoning_response is None:
+                self.console.print(
+                    "You have not initialized the task yet. Please perform the basic testing following `next` option.",
+                    style="bold red",
+                )
+                return
+            with self.console.status("[bold green] PentestGPT Thinking...") as status:
+                generation_response = self.test_generation_handler(
+                    self.step_reasoning_response
+                )
+
             self.console.print(
-                "You can follow the instructions below to complete the tasks.",
+                "Below are the further details.",
                 style="bold green",
             )
             self.console.print(generation_response + "\n")
             response = generation_response
 
-        # ask for sub tasks
-        elif request_option == "2":
+        # ask for task list (to-do list)
+        elif request_option == "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)
@@ -190,7 +217,7 @@ class pentestGPT:
             response = reasoning_response
 
         # pass other information, such as questions or some observations.
-        elif request_option == "3":
+        elif request_option == "discuss":
             ## (1) Request for user multi-line input
             self.console.print("Please share your thoughts/questions with PentestGPT.")
             user_input = prompt_ask(
@@ -204,7 +231,7 @@ class pentestGPT:
             self.console.print(response + "\n", style="yellow")
 
         # end
-        elif request_option == "4":
+        elif request_option == "quit":
             response = False
             self.console.print("Thank you for using PentestGPT!", style="bold green")
 
@@ -249,6 +276,9 @@ 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
                 break
 

utils/task_handler.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+"""
+url: https://github.com/prompt-toolkit/python-prompt-toolkit/tree/master/examples/prompts/auto-completion
+Demonstration of a custom completer class and the possibility of styling
+completions independently by passing formatted text objects to the "display"
+and "display_meta" arguments of "Completion".
+"""
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.formatted_text import HTML
+from prompt_toolkit.shortcuts import CompleteStyle, prompt
+
+
+class mainTaskCompleter(Completer):
+    tasks = [
+        "next",
+        "more",
+        "todo",
+        "discuss",
+        "help",
+        "quit",
+    ]
+
+    task_meta = {
+        "next": HTML("Go to the next step."),
+        "more": HTML("Explain the task with more details."),
+        "todo": HTML("Ask <b>PentestGPT</b> for todos."),
+        "discuss": HTML("Discuss with <b>PentestGPT</b>."),
+        "help": HTML("Show the help page."),
+        "quit": HTML("End the current session."),
+    }
+
+    task_details = """
+Below are the available tasks:
+ - next: Continue to the next step by inputting the test results.
+ - more: Explain the previous given task with more details.
+ - todo: Ask PentestGPT for the task list and what to do next.
+ - discuss: Discuss with PentestGPT. You can ask for help, discuss the task, or give any feedbacks.
+ - help: Show this help page.
+ - quit: End the current session."""
+
+    def get_completions(self, document, complete_event):
+        word = document.get_word_before_cursor()
+        for task in self.tasks:
+            if task.startswith(word):
+                yield Completion(
+                    task,
+                    start_position=-len(word),
+                    display=task,
+                    display_meta=self.task_meta.get(task),
+                )
+
+
+def main_task_entry(text="> "):
+    """
+    Entry point for the task prompt. Auto-complete
+    """
+    task_completer = mainTaskCompleter()
+    while True:
+        result = prompt(text, completer=task_completer)
+        if result not in task_completer.tasks:
+            print("Invalid task, try again.")
+        else:
+            return result
+
+
+if __name__ == "__main__":
+    main_task_entry()