feat: 🎸 Input length parser

✅ Closes: #2

README.md

@@ -1,5 +1,5 @@
 # PentestGPT
-v0.1.1, 09/04/2023
+v0.1, 09/04/2023
 
 ## Introduction
 **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.
@@ -16,8 +16,8 @@ The project is still in its early stage. Feel free to raise any issues when usin
 
 
 ## Examples
-2. To start, run `python3 main.py`. 
-3. The tool works similar to *msfconsole*. Follow the guidance to perform penetration testing. 
+1. To start, run `python3 main.py`. 
+2. The tool works similar to *msfconsole*. Follow the guidance to perform penetration testing. 
 
 ## Development
 - [ ] Add chunk processing

utils/pentest_gpt.py

@@ -7,7 +7,7 @@ from rich.console import Console
 
 
 import loguru
-import time, os
+import time, os, textwrap
 
 
 logger = loguru.logger
@@ -86,7 +86,21 @@ Do you understand?"""
                 logger.error(e)
 
     def _ask(self, text="> ", multiline=True) -> str:
-        # a handler for Prompt.ask that can intake multiple lines
+        """
+        A handler for Prompt.ask. It can intake multiple lines. Ideally for tool outputs and web contents
+
+        Parameters
+        ----------
+        text : str, optional
+            The prompt text, by default "> "
+        multiline : bool, optional
+            Whether to allow multiline input, by default True
+
+        Returns
+        -------
+        str
+            The user input
+        """
         if not multiline:
             return self.console.input(text)
         response = [self.console.input(text)]
@@ -111,10 +125,20 @@ Do you understand?"""
         # do some engineering trick here. Add postfix to the input to make it more understandable by LLMs.
         if source is not None and source in self.postfix_options.keys():
             prefix = prefix + self.postfix_options[source]
-        # TODO: chunk processing
-        summarized_content = self.chatGPTAgent.send_message(
-            prefix + text, self.input_parsing_session_id
-        )
+        # The default token-size limit is 4096 (web UI even shorter). 1 token ~= 4 chars in English
+        # Use textwrap to split inputs. Limit to 2000 token (8000 chars) for each input
+        # (1) replace all the newlines with spaces
+        text = text.replace("\r", " ").replace("\n", " ")
+        # (2) wrap the text
+        wrapped_text = textwrap.fill(text, 8000)
+        wrapped_inputs = wrapped_text.split("\n")
+        # (3) send the inputs to chatGPT input_parsing_session and obtain the results
+        summarized_content = ""
+        for wrapped_input in wrapped_inputs:
+            word_limit = f"Please ensure that the input is less than {8000 / len(wrapped_input)} words.\n"
+            summarized_content += self.chatGPTAgent.send_message(
+                prefix + word_limit + text, self.input_parsing_session_id
+            )
         return summarized_content
 
     def test_generation_handler(self):
@@ -182,9 +206,13 @@ Do you understand?"""
 
         # pass other information, such as questions or some observations.
         elif request_option == "3":
-            ## (1) pass the information to the input_parsing session.
-            response = Prompt.ask("> ")
-            ## (2) pass the summarized information to the reasoning session.
+            ## (1) Request for user multi-line input
+            self.console.print("Please input your information. End with EOF.")
+            user_input = self._ask("> ", multiline=True)
+            ## (2) directly pass the information to the reasoning session.
+            prefix = "The tester provides the following thoughts for your consideration. Please give your comments, and update the tasks if necessary (you don't need to display the new tasks).\n"
+            response = self.reasoning_handler(prefix + user_input)
+
         # end
         elif request_option == "4":
             response = False