update black

README.md

@@ -10,4 +10,4 @@ The goal is to build a user-friendly bridge to automate ChatGPT and other GPT mo
 
 ## Examples
 1. **Get username on local PC**: `python3 example_chatgpt_api.py`
-2. **Use SQLMap**: `python3 example_sqlmap.py`
+2. **Use SQLMap**: `python3 example_sqlmap.py`

example_chatgpt_api.py

@@ -3,9 +3,12 @@ from llm_handle.parser import extract_cmd
 from task_handle.cmd_execution import execute_cmd
 import os
 
+
 def __main__():
     bot = ChatGPT()
-    response = bot.ask("Can you give me a sample command in Mac terminal for checking the user names? Please give me the code directly.")
+    response = bot.ask(
+        "Can you give me a sample command in Mac terminal for checking the user names? Please give me the code directly."
+    )
     sample_response = """
     Certainly! To list all user names on a Mac using the terminal, you can use the `dscl` command with the `list` option for the `/Users` node. Here's the command:
     ```
@@ -30,6 +33,3 @@ def __main__():
 
     # delete the session in the end
     bot.delete_conversation()
-
-
-

example_sqlmap.py

@@ -21,11 +21,14 @@ Are you clear about it?
 keyword = "vulnerability detected!!!"
 prefix = "The output from terminal is :\n"
 
+
 class sqlmapHandler(chatGPTTemplate):
     # should override the run function
     def run(self):
         self.initialize()
-        response = self.ask("Now please start, the website is: http://testphp.vulnweb.com/listproducts.php?cat=1")
+        response = self.ask(
+            "Now please start, the website is: http://testphp.vulnweb.com/listproducts.php?cat=1"
+        )
         while True:
             # get the response from the bot
             # if the keyword is detected, break the loop
@@ -47,11 +50,12 @@ class sqlmapHandler(chatGPTTemplate):
                 # feed the output to the bot
                 response = self.ask(output, need_prefix=True)
 
+
 if __name__ == "__main__":
-    #1. init the bot session
+    # 1. init the bot session
     bot = ChatGPT()
     chat_handler = sqlmapHandler(bot, init_script=init_script)
     chat_handler._update_prefix(prefix)
 
-    #2. run the chat
+    # 2. run the chat
     chat_handler.run()

llm_handle/parser.py

@@ -1,8 +1,10 @@
 import re
-def extract_cmd(response:str) -> str:
+
+
+def extract_cmd(response: str) -> str:
     """
     Process the response from chatgpt_wrapper, and extract the command for the bot.
-    
+
     Parameters
     ----------
     response: str
@@ -27,11 +29,9 @@ def extract_cmd(response:str) -> str:
             raise ValueError("More than one command is found.")
     except AttributeError:  # Nonetype, nothing found
         return False
-    
+
     result = result_list[0]
-    if result[0] == "\n": # If the command starts with a newline, remove it.
+    if result[0] == "\n":  # If the command starts with a newline, remove it.
         result = result[1:]
-    
+
     return result
-
-

requirements.txt

@@ -2,3 +2,4 @@ requests
 pyyaml
 playwright==1.28.0
 sqlmap
+black

task_handle/cmd_execution.py

@@ -1,15 +1,15 @@
-
 import os, subprocess
 
-def execute_cmd(cmd:str) -> str:
+
+def execute_cmd(cmd: str) -> str:
     """
     Execute the command in the mac terminal.
-    
+
     Parameters
     ----------
     cmd: str
         The command to be executed.
-    
+
     Returns
     ----------
     output: str
@@ -17,17 +17,17 @@ def execute_cmd(cmd:str) -> str:
     """
     try:
         # execute the command in the system terminal
-        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr = None, shell=True)
+        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=None, shell=True)
         output = ""
         # some tools may take time to execute. Wait until the output is finished.
         while True:
             line_output = p.stdout.readline()
             if line_output:
                 output += line_output.decode("utf-8")
-            if line_output == b'' and p.poll() is not None:
+            if line_output == b"" and p.poll() is not None:
                 break
 
         return output
     except Exception as e:
         print("Error in executing the command:", e)
-        return None
+        return None

task_handle/custom_exceptions.py

@@ -1,4 +1,5 @@
 # declare the custom types of exceptions
 
+
 class NoCodeFromResponseException(Exception):
-    pass
+    pass

task_handle/template.py

@@ -4,7 +4,7 @@ import os, logging, re
 from .custom_exceptions import NoCodeFromResponseException
 
 
-class chatGPTTemplate():
+class chatGPTTemplate:
     """
     A template for the chatGPT task.
     It contains the basic functions that are required for the task.
@@ -13,7 +13,7 @@ class chatGPTTemplate():
     def __init__(self, bot_session, init_script=None):
         """
         Initialize the by taking the session
-        The bot session is a standard chatgpt_wrapper bot session. 
+        The bot session is a standard chatgpt_wrapper bot session.
         More details at https://github.com/mmabrouk/chatgpt-wrapper
 
         Parameters:
@@ -34,7 +34,7 @@ class chatGPTTemplate():
         self._init_script = init_script
         self._prefix = None
         self._exception_ask = {}
-    
+
     def _extract_command(self, response: str) -> str:
         """
         This function is used to extract the command from the response.
@@ -55,11 +55,11 @@ class chatGPTTemplate():
                 raise ValueError("More than one command is found.")
         except Exception:  # Nonetype, nothing found
             raise NoCodeFromResponseException("No code is found in the response.")
-        
+
         result = result_list[0]
-        if result[0] == "\n": # If the command starts with a newline, remove it.
+        if result[0] == "\n":  # If the command starts with a newline, remove it.
             result = result[1:]
-        
+
         return result
 
     def _cmd_wrapper(self, cmd: str) -> str:
@@ -73,8 +73,7 @@ class chatGPTTemplate():
         # the possible types of exceptions
         output = execute_cmd(cmd)
         return output
-        
-    
+
     def _update_init_script(self, init_script: str):
         """
         This function is used to update the initialization script.
@@ -101,7 +100,6 @@ class chatGPTTemplate():
         Returns:
         """
         return prefix + question
-    
 
     ########## Implementations ##########
 
@@ -112,8 +110,8 @@ class chatGPTTemplate():
         """
         if self._init_script is not None:
             self._bot_session.ask(self._init_script)
-    
-    def ask(self, question: str, need_prefix = False) -> str:
+
+    def ask(self, question: str, need_prefix=False) -> str:
         """
         Wrap the default bot ask function.
         Parameters:
@@ -133,7 +131,7 @@ class chatGPTTemplate():
         except Exception as e:
             print("Error in asking the question:", e)
             return None
-    
+
     def exception_ask(self, question: str) -> str:
         """
         This function is used to ask the bot when an exception is raised.
@@ -146,10 +144,10 @@ class chatGPTTemplate():
             return self.ask(self._exception_ask)
         else:
             return None
-    
+
     def run(self):
         """
         The function with the main logic. This should be overwritten in the task execution.
         """
         print("Please override the run function!")
-        pass
+        pass

tasks/NLP_RESTAPI_Repairment.py

@@ -1,10 +1,5 @@
 # This handle helps to translate RESTful API documentation in natural language to OpenAPI Specification 3.0
 
 
-
 ##### Functions #####
 # Get the natural language description of the RESTful API.
-
-
-
-

test_os_execution.py

@@ -6,16 +6,15 @@ import subprocess
 cmd = 'sqlmap -u "http://testphp.vulnweb.com/listproducts.php?cat=1" --batch --level=5 --risk=3'
 
 # execute the command
-p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr = None, shell=True)
+p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=None, shell=True)
 output_str = ""
 while True:
     output = p.stdout.readline()
     if output:
         print(output.decode("utf-8"), end="")
         output_str += output.decode("utf-8")
-    if output == b'' and p.poll() is not None:
+    if output == b"" and p.poll() is not None:
         print("------end of output------")
         break
 
 print(output_str)
-