diff --git a/README.md b/README.md index 00b7c96..2d78c76 100644 --- 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. @@ -45,14 +44,28 @@ https://user-images.githubusercontent.com/78410652/232327920-7318a0c4-bee0-4cb4- - Fill in `userAgent` with your user agent. - 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 + ## Test connection for OpenAI api (GPT-4) + 2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use + ## 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 + ``` +4. (Notice) The above verification process for cookie. 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`. Please submit an issue if you encounter any problem. ## 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: @@ -72,7 +85,7 @@ https://user-images.githubusercontent.com/78410652/232327920-7318a0c4-bee0-4cb4- - `brainstorm`: let PentestGPT brainstorm on the local task for all the possible solutions. - `discuss`: discuss with PentestGPT about this local task. - `google`: search on Google. This function is still under development. - - `continue`: exit the sub task and continue the main testing session. + - `continue`: exit the subtask and continue the main testing session. ## Report 1. After finishing the penetration testing, a report will be automatically generated in `logs` folder (if you quit with `quit` command). 2. The report can be printed in a human-readable format by running `python3 utils/report_generator.py `. A sample report `sample_pentestGPT_log.txt` is also uploaded. diff --git a/main.py b/main.py index 76f8b76..d8414a3 100644 --- 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 index 286b141..8410899 100644 --- a/test_connection.py +++ b/test_connection.py @@ -2,27 +2,66 @@ import loguru 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 " + "1. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use " ) 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 ") + 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 ") + 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" + ) + +