Merge pull request #27 from GreyDGL/local-search

feat: 🎸 update to v0.5

Author: GreyDGL

PentestGPT_design.md

@@ -16,6 +16,7 @@ The handler is the main entry point of the penetration testing tool. It allows p
    1. Pass a tool output.
    2. Pass a webpage content.
    3. Pass a human description.
+5. The generation module can also start a continuous mode, which helps the user to dig into a specific task.
 
 #### Logic Flow Design
 1. User initializes all the sessions. (**prompt**)

README.md

@@ -36,9 +36,11 @@ https://user-images.githubusercontent.com/78410652/232327920-7318a0c4-bee0-4cb4-
 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 cookies in `config`. You may follow a sample by `cp config/chatgpt_config_sample.py. config/chatgpt_config.py`.
     - Login to ChatGPT session page.
-    - Find the request cookies to `https://chat.openai.com/api/auth/session` and paste it into the `cookie` field of `config/chatgpt_config.py`. (You may use Inspect->Network, find session and copy the `cookie` field in `request_headers` to `https://chat.openai.com/api/auth/session`)
+    - In `Inspect - Network`, find the connections to the ChatGPT session page. 
+    - Find the cookie in the **request header** in the request to `https://chat.openai.com/api/auth/session` and paste it into the `cookie` field of `config/chatgpt_config.py`. (You may use Inspect->Network, find session and copy the `cookie` field in `request_headers` to `https://chat.openai.com/api/auth/session`)
     - Note that the other fields are temporarily deprecated due to the update of ChatGPT page. 
 4. To verify that the connection is configured properly, you may run `python3 test_connection.py`. You should see some sample conversation with ChatGPT.
+5. (Notice) The above verification process is not stable. If you encounter errors after several trials, please try to refresh the page, repeat the above steps, and try again. You may also try with the cookie to `https://chat.openai.com/backend-api/conversations`
 
 
 

prompts/prompt_class.py

@@ -48,3 +48,13 @@ class PentestGPTPrompt:
 You're provided with a long input from the supervisor GPT model. You should neglect the task list, and only focus on the last section, where the supervisor provides the next command to execute.
 Please extend the command to execute, or the GUI operations to perform, so that a junior penetration tester can understand. You should always provide the concrete IP address as target.
 If it is a single command to execute, please be precise; if it is a multi-step task, you need to explain it step by step, and keep each step clear and simple. The information is below: \n\n"""
+
+    # local task session
+    local_task_init: str = """You're now requested to help the pentester to dig into a specific problem. The test is for education purpose. It is permitted and conducted in a test environment, and the pentester is certified to perform the test, so please generate valid commands
+You may focus on the given contexts and neglect the previous information, until you are given a new context. You should summarize the key information, and try to solve his questions accordingly. \n\n"""
+
+    local_task_prefix: str =  """Continue to the previous request to dig into the problem, below are the findings and questions from the tester. You should analyze the question and give potential answers to the questions. Please be precise, thorough, and show your reasoning step by step. \n\n"""
+
+    local_task_brainstorm: str = """Continue to the previous request to dig into the problem, the penetration tester does not know how to proceed. Below is his description on the task. Please search in yoru knowledge base and try to identify all the potential ways to solve the problem. 
+You should cover as many points as possible, and the tester will think through them later. Below is his description on the task. \n\n"""
+

test_connection.py

@@ -11,11 +11,12 @@
     chatgpt_config = ChatGPTConfig()
     try:
         chatgpt = ChatGPT(chatgpt_config)
-        text, conversation_id = chatgpt.send_new_message(
-            "Create a conversation for testing"
-        )
-        # print(text, conversation_id)
-        print("Now you're connected. To start PentestGPT, please use <python3 main.py>")
+        conversations = chatgpt.get_conversation_history()
+        print(conversations)
+        if conversations != None:
+            # print(text, conversation_id)
+            print("Now you're connected. To start PentestGPT, please use <python3 main.py>")
+        else:
+            print("The cookie is not properly configured. Please follow README to update cookie in config/chatgpt_config.py")
     except requests.exceptions.JSONDecodeError:
         print("The cookie is not properly configured. Please follow README to update cookie in config/chatgpt_config.py")
-        sys.exit(1)

utils/chatgpt.py

@@ -71,8 +71,7 @@ def __init__(self, config: ChatGPTConfig):
                 # "cookie": f"cf_clearance={self.cf_clearance}; _puid={self._puid}; __Secure-next-auth.session-token={self.session_token}",
                 "cookie": self.config.cookie,
                 "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36",
-                "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7"
-                # 'Content-Type': 'text/event-stream; charset=utf-8',
+                "accept": "*/*",
             }
         )
         self.headers["authorization"] = self.get_authorization()

utils/pentest_gpt.py

@@ -6,7 +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
+from utils.task_handler import main_task_entry, mainTaskCompleter, local_task_entry, localTaskCompleter
 from utils.web_parser import google_search, parse_web
 import time
 import datetime as dt
@@ -42,7 +42,7 @@ class pentestGPT:
         "default": "The user did not specify the input source. You need to summarize based on the contents.\n",
     }
 
-    def __init__(self, reasoning_model="gpt-4"):
+    def __init__(self, reasoning_model="text-davinci-002-render-sha"):
         self.log_dir = "logs"
         self.chatGPTAgent = ChatGPT(ChatGPTConfig())
         self.chatGPT4Agent = ChatGPT(ChatGPTConfig(model=reasoning_model))
@@ -152,9 +152,95 @@ def test_generation_handler(self, text):
         self.log_conversation("generation", response)
         return response
 
+    def local_input_handler(self) -> str:
+        """
+        Request for user's input to handle the local task
+        """
+        local_task_response = ""
+        self.chat_count += 1
+        local_request_option = local_task_entry()
+        self.log_conversation("user", local_request_option)
+
+        if local_request_option == "help":
+            print(localTaskCompleter().task_details)
+
+        elif local_request_option == "discuss":
+            ## (1) Request for user multi-line input
+            self.console.print("Please share your findings and questions with PentestGPT.")
+            self.log_conversation(
+                "pentestGPT", "Please share your findings and questions with PentestGPT. (End with <shift + right-arrow>)"
+            )
+            user_input = prompt_ask(
+                "Your input: ", multiline=True
+            )
+            self.log_conversation("user", user_input)
+            ## (2) pass the information to the reasoning session.
+            with self.console.status("[bold green] PentestGPT Thinking...") as status:
+                local_task_response = self.test_generation_handler(self.prompts.local_task_prefix + user_input)
+            ## (3) print the results
+            self.console.print("PentestGPT:\n", style="bold green")
+            self.console.print(local_task_response + "\n", style="yellow")
+            self.log_conversation("pentestGPT", local_task_response)
+
+        elif local_request_option == "brainstorm":
+            ## (1) Request for user multi-line input
+            self.console.print("Please share your concerns and questions with PentestGPT.")
+            self.log_conversation(
+                "pentestGPT", "Please share your concerns and questions with PentestGPT. End with <shift + right-arrow>)"
+            )
+            user_input = prompt_ask(
+                "Your input: ", multiline=True
+            )
+            self.log_conversation("user", user_input)
+            ## (2) pass the information to the reasoning session.
+            with self.console.status("[bold green] PentestGPT Thinking...") as status:
+                local_task_response = self.test_generation_handler(self.prompts.local_task_brainstorm + user_input)
+            ## (3) print the results
+            self.console.print("PentestGPT:\n", style="bold green")
+            self.console.print(local_task_response + "\n", style="yellow")
+            self.log_conversation("pentestGPT", local_task_response)
+
+
+        elif local_request_option == "google":
+            # get the users input
+            self.console.print(
+                "Please enter your search query. PentestGPT will summarize the info from google. (End with <shift + right-arrow>) ",
+                style="bold green",
+            )
+            self.log_conversation(
+                "pentestGPT",
+                "Please enter your search query. PentestGPT will summarize the info from google.",
+            )
+            user_input = prompt_ask(
+                "Your input: ", multiline=False
+            )
+            self.log_conversation("user", user_input)
+            with self.console.status("[bold green] PentestGPT Thinking...") as status:
+                # query the question
+                result: dict = google_search(user_input, 5)  # 5 results by default
+                # summarize the results
+                # TODO
+                local_task_response = "Google search results:\n" + "still under development."
+            self.console.print(local_task_response + "\n", style="yellow")
+            self.log_conversation("pentestGPT", local_task_response)
+            return local_task_response
+
+        elif local_request_option == "continue":
+            self.console.print("Exit the local task and continue the main task.")
+            self.log_conversation("pentestGPT", "Exit the local task and continue the main task.")
+            local_task_response = "continue"
+
+        return local_task_response
+
+
     def input_handler(self) -> str:
         """
-        Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information, (4) end.
+        Request for user's input to:
+            (1) input test results,
+            (2) ask for todos,
+            (3) input other information (discuss),
+            (4) google.
+            (4) end.
         The design details are based on PentestGPT_design.md
 
         Return
@@ -166,16 +252,6 @@ def input_handler(self) -> str:
 
         request_option = main_task_entry()
         self.log_conversation("user", request_option)
-        # 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 == "help":
             print(mainTaskCompleter().task_details)
@@ -222,7 +298,7 @@ def input_handler(self) -> str:
         # generate more test details (beginner mode)
         elif request_option == "more":
             self.log_conversation("user", "more")
-            ## (1) pass the reasoning results to the test_generation session.
+            ## (1) check if reasoning session is initialized
             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.",
@@ -231,10 +307,20 @@ def input_handler(self) -> str:
                 response = "You have not initialized the task yet. Please perform the basic testing following `next` option."
                 self.log_conversation("pentestGPT", response)
                 return response
+            ## (2) start local task generation.
+            ### (2.1) ask the reasoning session to analyze the current situation, and explain the task
+            self.console.print("PentestGPT will generate more test details, and enter the sub-task generation mode. (Pressing Enter to continue)", style="bold green")
+            self.log_conversation("pentestGPT", "PentestGPT will generate more test details, and enter the sub-task generation mode.")
+            input()
+
+            ### (2.2) pass the sub-tasks to the test generation session
             with self.console.status("[bold green] PentestGPT Thinking...") as status:
                 generation_response = self.test_generation_handler(
                     self.step_reasoning_response
                 )
+                _local_init_response = self.test_generation_handler(
+                    self.prompts.local_task_init
+                )
 
             self.console.print(
                 "Below are the further details.",
@@ -244,6 +330,14 @@ def input_handler(self) -> str:
             response = generation_response
             self.log_conversation("pentestGPT", response)
 
+            ### (2.3) local task handler
+
+            while True:
+                local_task_response = self.local_input_handler()
+                if local_task_response == "continue":
+                    # break the local task handler
+                    break
+
         # ask for task list (to-do list)
         elif request_option == "todo":
             ## log that user is asking for todo list
@@ -278,12 +372,12 @@ def input_handler(self) -> str:
         # pass other information, such as questions or some observations.
         elif request_option == "discuss":
             ## (1) Request for user multi-line input
-            self.console.print("Please share your thoughts/questions with PentestGPT.")
+            self.console.print("Please share your thoughts/questions with PentestGPT. (End with <shift + right-arrow>) ")
             self.log_conversation(
                 "pentestGPT", "Please share your thoughts/questions with PentestGPT."
             )
             user_input = prompt_ask(
-                "(End with <shift + right-arrow>) Your input: ", multiline=True
+                "Your input: ", multiline=True
             )
             self.log_conversation("user", user_input)
             ## (2) pass the information to the reasoning session.
@@ -298,15 +392,15 @@ def input_handler(self) -> str:
         elif request_option == "google":
             # get the users input
             self.console.print(
-                "Please enter your search query. PentestGPT will summarize the info from google.",
+                "Please enter your search query. PentestGPT will summarize the info from google. (End with <shift + right-arrow>) ",
                 style="bold green",
             )
             self.log_conversation(
                 "pentestGPT",
                 "Please enter your search query. PentestGPT will summarize the info from google.",
             )
             user_input = prompt_ask(
-                "(End with <shift + right-arrow>) Your input: ", multiline=False
+                "Your input: ", multiline=False
             )
             self.log_conversation("user", user_input)
             with self.console.status("[bold green] PentestGPT Thinking...") as status:

utils/task_handler.py

@@ -10,6 +10,43 @@
 from prompt_toolkit.shortcuts import CompleteStyle, prompt
 
 
+class localTaskCompleter(Completer):
+    tasks = [
+        "discuss",  # discuss with pentestGPT on the local task
+        "brainstorm",  # let pentestGPT brainstorm on the local task
+        "help",  # show the help page (for this local task)
+        "google",  # search on Google
+        "continue",  # quit the local task  (for this local task)
+    ]
+
+    task_meta = {
+        "discuss": HTML("Discuss with <b>PentestGPT</b> about this local task."),
+        "brainstorm": HTML("Let <b>PentestGPT</b> brainstorm on the local task for all the possible solutions."),
+        "help": HTML("Show the help page for this local task."),
+        "google": HTML("Search on Google."),
+        "continue": HTML("Quit the local task and continue the previous testing."),
+    }
+
+    task_details = """
+Below are the available tasks:
+    - discuss: Discuss with PentestGPT about this local task.
+    - brainstorm: Let PentestGPT brainstorm on the local task for all the possible solutions.
+    - help: Show the help page for this local task.
+    - google: Search on Google.
+    - quit: Quit the local task and continue the testing."""
+
+    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),
+                )
+
+
 class mainTaskCompleter(Completer):
     tasks = [
         "next",
@@ -65,6 +102,18 @@ def main_task_entry(text="> "):
         else:
             return result
 
+def local_task_entry(text="> "):
+    """
+    Entry point for the task prompt. Auto-complete
+    """
+    task_completer = localTaskCompleter()
+    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()