feat: 🎸 v0.3.0. Feature optimization

GreyDGL · Apr 14, 2023 · a151a55 · a151a55
1 parent 1bb5a08
commit a151a55
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 23 deletions.
diff --git a/README.md b/README.md
@@ -5,23 +5,36 @@ v0.2, 12/04/2023
 **PentestGPT** is a penetration testing tool empowered by **ChatGPT**. It is designed to automate the penetration testing process. It is built on top of ChatGPT and operate in an interactive mode to guide penetration testers in both overall progress and specific operations.
 A sample testing process of **PentestGPT** on a target VulnHub machine (Hackable II) is available at [here](./resources/PentestGPT_Hackable2.pdf).
 
+- Comparison to **Auto-GPT**. 
+  - Using [Auto-GPT](https://github.com/Torantulino/Auto-GPT) in security testing is good, but it is not optimized for security-related tasks. 
+  - **PentestGPT** is designed for penetration testing with a customized session interaction (see [here](./PentestGPT_design.md) for the detailed design).
+  - Currently, **PentestGPT** does not rely on search engine. The "Google-enhanced" version of **PentestGPT** is under development.
 
 ## Contribute
-The project is still in its early stage. Feel free to raise any issues when using the tool. 
+- The project is still in its early stage. Feel free to raise any issues when using the tool. 
+- This project is for research purpose. Please contact me if you're interested in collaboration.
 
 ## Installation
 1. Install `requirements.txt` with `pip install -r requirements.txt`
-2. Install `chatgpt-wrapper` if you're non-plus members: `pip install git+https://github.com/mmabrouk/chatgpt-wrapper`. More details at: https://github.com/mmabrouk/chatgpt-wrapper. Note that the support for non-plus members are not optimized.
+2. (Deprecated: Will update support for non-plus member later.) ~~Install `chatgpt-wrapper` if you're non-plus members: `pip install git+https://github.com/mmabrouk/chatgpt-wrapper`. More details at: https://github.com/mmabrouk/chatgpt-wrapper. Note that the support for non-plus members are not optimized.~~
 3. Configure the keys in `config`. You may follow a sample by `cp config/chatgpt_config_sample.py. config/chatgpt_config.py`.
 
 
 
 ## Usage
 1. To start, run `python3 main.py`. 
 2. The tool works similar to *msfconsole*. Follow the guidance to perform penetration testing. 
-3. In general, PentestGPT intakes commands similar to chatGPT. 
-   - To intake multi-line inputs in the terminal, please use <Enter> for new line, and <Shift+Right-Arror> to submit the input.
-   - The selection bar allows you to select a pre-defined options.
+3. In general, PentestGPT intakes commands similar to chatGPT. There are several basic commands.
+   1. The commands are: 
+      - `help`: show the help message.
+      - `next`: key in the test execution result and get the next step.
+      - `more`: let **PentestGPT** to explain more details of the current step.
+      - `todo`: show the todo list.
+      - `discuss`: discuss with the **PentestGPT**.
+      - `exit`: exit the tool.
+   2. You can use <SHIFT + right arrow> to end your input (and <ENTER> is for next line).
+   3. You may always use `TAB` to autocomplete the commands.
+   4. When you're given a drop-down selection list, you can use cursor or arrow key to navigate the list. Press `ENTER` to select the item. Similarly, use <SHIFT + right arrow> to confirm selection.
 
 
 ## Design Documentation
@@ -44,6 +57,9 @@ The handler is the main entry point of the penetration testing tool. It allows p
    3. Pass a human description.
 
 ## Update history
+### v0.3
+- Prompt usage optimization.
+- Documentation improvements.
 ### v0.2
 - A major update to improve the terminal usage
 - Prompt optimization.
diff --git a/utils/pentest_gpt.py b/utils/pentest_gpt.py
@@ -6,6 +6,7 @@
 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 @@ def __init__(self):
         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 @@ def input_handler(self) -> str:
         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 @@ def input_handler(self) -> str:
                 )
                 ## (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 @@ def input_handler(self) -> str:
             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 @@ def input_handler(self) -> str:
             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 @@ def main(self):
         # 4. enter the main loop.
         while True:
             result = self.input_handler()
+            self.console.print(
+                "-----------------------------------------", style="bold white"
+            )
             if not result:  # end the session
                 break
 

diff --git a/utils/task_handler.py b/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()