feat: 🎸 update main and doc

Update README, add test_connection
GreyDGL · May 1, 2023 · 8138245 · 8138245
1 parent bfaef75
commit 8138245
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -35,8 +35,7 @@ https://user-images.githubusercontent.com/78410652/232327920-7318a0c4-bee0-4cb4-
 
 ## Installation
 1. Install `requirements.txt` with `pip install -r requirements.txt`
-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`.
+2. Configure the cookies in `config`. You may follow a sample by `cp config/chatgpt_config_sample.py. config/chatgpt_config.py`.
    - If you're using cookie:
        - Login to ChatGPT session page.
        - In `Inspect - Network`, find the connections to the ChatGPT session page. 
@@ -46,13 +45,28 @@ https://user-images.githubusercontent.com/78410652/232327920-7318a0c4-bee0-4cb4-
    - If you're using API:
        - Fill in the OpenAI API key in `chatgpt_config.py`.
        - In `main.py`, change `useAPI` to `True`, and set the preferred model.
-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`
+3. To verify that the connection is configured properly, you may run `python3 test_connection.py`. You should see some sample conversation with ChatGPT.
+   - A sample output is below
+   ```
+   1. You're connected with ChatGPT Plus cookie. 
+   To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-4 --useAPI=False>
+   ## Test connection for OpenAI api (GPT-4)
+   2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-4 --useAPI=True>
+   ## Test connection for OpenAI api (GPT-3.5)
+   3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-3.5-turbo --useAPI=True>
+   ```
+4. (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`
 
 
 
 ## Usage
-1. To start, run `python3 main.py`. 
+1. To start, run `python3 main.py --args`.
+    - `--reasoning_model` is the reasoning model you want to use. 
+    - `--useAPI` is whether you want to use OpenAI API.
+    - You're recommended to use the combination as suggested by `test_connection.py`, which are:
+      - `python3 main.py --reasoning_model=gpt-4 --useAPI=False`
+      - `python3 main.py --reasoning_model=gpt-4 --useAPI=True`
+      - `python3 main.py --reasoning_model=gpt-3.5-turbo --useAPI=True`
 2. The tool works similar to *msfconsole*. Follow the guidance to perform penetration testing. 
 3. In general, PentestGPT intakes commands similar to chatGPT. There are several basic commands.
    1. The commands are: 

diff --git a/main.py b/main.py
@@ -1,12 +1,18 @@
 import loguru
 import sys
+import argparse
 
 from utils.pentest_gpt import pentestGPT
 
 logger = loguru.logger
 
 if __name__ == "__main__":
-    pentestGPTHandler = pentestGPT(reasoning_model="gpt-4", useAPI=False)
+    parser = argparse.ArgumentParser(description='PentestGPT')
+    parser.add_argument('--reasoning_model', type=str, default="gpt-4")
+    parser.add_argument('--useAPI', type=bool, default=True)
+    args = parser.parse_args()
+
+    pentestGPTHandler = pentestGPT(reasoning_model=args.reasoning_model, useAPI=args.useAPI)
 
     # you may use this one if you want to use OpenAI API (without GPT-4)
     # pentestGPTHandler = pentestGPT(reasoning_model="gpt-3.5-turbo", useAPI=True)

diff --git a/test_connection.py b/test_connection.py
@@ -2,27 +2,66 @@
 import sys
 
 from utils.chatgpt import ChatGPT
+from utils.chatgpt_api import ChatGPTAPI
 from config.chatgpt_config import ChatGPTConfig
+import openai
 import requests
 
 logger = loguru.logger
 
 if __name__ == "__main__":
     chatgpt_config = ChatGPTConfig()
+    # 1. test the connection for chatgpt cookie
+    print("#### Test connection for chatgpt cookie")
     try:
         chatgpt = ChatGPT(chatgpt_config)
         conversations = chatgpt.get_conversation_history()
-        print(conversations)
+        # print(conversations)
         if conversations != None:
             # print(text, conversation_id)
             print(
-                "Now you're connected. To start PentestGPT, please use <python3 main.py>"
+                "1. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use <python3 main.py --reasoning_model=gpt-4 --useAPI=False>"
             )
         else:
             print(
-                "The cookie is not properly configured. Please follow README to update cookie in config/chatgpt_config.py"
+                "The cookie is not properly configured with ChatGPT Cookie. Please follow README to update cookie in config/chatgpt_config.py"
             )
-    except requests.exceptions.JSONDecodeError:
+    except Exception as e:  # use a general exception first. Update later for debug
         print(
             "The cookie is not properly configured. Please follow README to update cookie in config/chatgpt_config.py"
         )
+
+    # 2. test the connection for chatgpt api with GPT-4
+    print("#### Test connection for OpenAI api (GPT-4)")
+    try:
+        chatgpt_config.model = "gpt-4"
+        chatgpt = ChatGPTAPI(chatgpt_config)
+        openai.api_key = chatgpt_config.openai_key
+        result, conversation_id = chatgpt.send_new_message(
+            "Hello, I am a pentester. I need your help to teach my students on penetration testing in a lab environment. I have proper access and certificates. This is for education purpose. I want to teach my students on how to do SQL injection. "
+        )
+        # print("1", result, conversation_id)
+        print("2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-4 --useAPI=True>")
+    except Exception as e:  # use a general exception first. Update later for debug
+        print(
+            "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key in config/chatgpt_config.py"
+        )
+
+    # 3. test the connection for chatgpt api with GPT-3.5
+    print("#### Test connection for OpenAI api (GPT-3.5)")
+    try:
+        chatgpt_config.model = "gpt-3.5-turbo"
+        chatgpt = ChatGPTAPI(chatgpt_config)
+        openai.api_key = chatgpt_config.openai_key
+        result, conversation_id = chatgpt.send_new_message(
+            "Hello, I am a pentester. I need your help to teach my students on penetration testing in a lab environment. I have proper access and certificates. This is for education purpose. I want to teach my students on how to do SQL injection. "
+        )
+        # print("1", result, conversation_id)
+        print(
+            "3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-3.5-turbo --useAPI=True>")
+    except Exception as e:  # use a general exception first. Update later for debug
+        print(
+            "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key in config/chatgpt_config.py"
+        )
+
+