weyne/PentestGPT
1
0
Fork 0
mirror of https://github.com/weyne85/PentestGPT.git synced 2025-10-29 16:58:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: 🎸 v0.1 publish

Browse Source
This commit is contained in:
Grey_D
2023-04-09 01:03:59 +08:00
parent e8bd7a6d5f
commit 797370edac
5 changed files with 257 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
PentestGPT_design.md
View File

@@ -1,13 +1,12 @@
# Design Documentation for PentestGPT
## Design Documentation for PentestGPT
version 0.1, for web penetration testing only
## General Design
### General Design
PentestGPT provides a unified terminal input handler, and backed by three main components:
- A test generation module which generates the exact penetration testing commands or operations for the users to execute.
- A test reasoning module which conducts the reasoning of the test, guiding the penetration testers on what to do next.
- A parsing module which parses the output of the penetration tools and the contents on the webUI.
## Function Design
### Function Design
The handler is the main entry point of the penetration testing tool. It allows pentesters to perform the following operations:
1. (initialize itself with some pre-designed prompts.)
2. Start a new penetration testing session by providing the target information.
@@ -18,14 +17,14 @@ The handler is the main entry point of the penetration testing tool. It allows p
3. Pass a human description.
## System Design
### General Structure
### System Design
#### General Structure
1. Maintain three chat sessions in one class. Each session is for one component.
2. User can select to pass information to one section. In particular.
1. todo.
2. pass information.
### Logic Flow Design
#### Logic Flow Design
1. User initializes all the sessions.
2. User follows the instruction to perform penetration testing in the following logic.
1. Reasoning session provides a list of todo for user to choose from.
@@ -62,9 +61,6 @@ sequenceDiagram
ParsingSession-->>User: 2.2 Summarize information
User-->>ReasoningSession: 2.3 Pass summarized information
```
### Wrapper Design
### Session Design
#### Wrapper Design
#### Session Design

92
README.md
View File

@@ -1,16 +1,94 @@
# GPT-Terminal-Bridge
# PentestGPT
## General
The goal is to build a user-friendly bridge to automate ChatGPT and other GPT modules to automatically control terminal-based applications. This allows an automation of many human-intensive works, such as penetration testing with tools.
## 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.
There are two modes of running. One is based on the project `chatgpt-wrapper`, which uses playwright browser to extend the ChatGPT API. It is recommended for free users because other API libraries cannot bypass the Cloudflare protection. The other mode is API-based, which utilizes `utils/chatgpt.py`. ChatGPT plus user can provide necessary key information to use this mode.
## Contribute
The project is still in its early stage. Feel free to raise any issues when using the tool.
## Installation
1. Install `requirements.txt` with `pip install -r requirements.txt`
2. Install `chatgpt-wrapper`: `pip install git+https://github.com/mmabrouk/chatgpt-wrapper`. More details at: https://github.com/mmabrouk/chatgpt-wrapper.
2. 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.
## Examples
1. **Test Example with ChatGPT Browser**: `python3 example_chatgpt_playwright.py`
2. **Test Example with ChatGPT API** (requires API key): `python3 example_chatgpt_API.py`
1. To start, run `python3 main.py`.
## Development
- [ ] Add chunk processing
- [ ] Add prompt optimization
- [ ] Test scenarios beyond web testing
## Design Documentation
The current design is mainly for web penetration testing
### General Design
PentestGPT provides a unified terminal input handler, and backed by three main components:
- A test generation module which generates the exact penetration testing commands or operations for the users to execute.
- A test reasoning module which conducts the reasoning of the test, guiding the penetration testers on what to do next.
- A parsing module which parses the output of the penetration tools and the contents on the webUI.
### Function Design
The handler is the main entry point of the penetration testing tool. It allows pentesters to perform the following operations:
1. (initialize itself with some pre-designed prompts.)
2. Start a new penetration testing session by providing the target information.
3. Ask for todo-list, and acquire the next step to perform.
4. After completing the operation, pass the information to PentestGPT.
1. Pass a tool output.
2. Pass a webpage content.
3. Pass a human description.
### System Design
#### General Structure
1. Maintain three chat sessions in one class. Each session is for one component.
2. User can select to pass information to one section. In particular.
1. todo.
2. pass information.
#### Logic Flow Design
1. User initializes all the sessions.
2. User follows the instruction to perform penetration testing in the following logic.
1. Reasoning session provides a list of todo for user to choose from.
2. User chooses one todo, and reasoning session passes the information to generation session.
3. Generation session generates the exact command for the user to execute.
4. User executes the command, and passes the output to the parsing session.
5. Parsing session parses the output, and passes the information to reasoning session.
6. Reasoning session updates the todo list, and provides the next todo for the user to choose from.
3. User can also pass information to the reasoning session directly.
1. User passes the information to parser session.
2. Parser session summarizes the information, and passes the information to the reasoning section.
A flow-chart is shown below:
```mermaid
sequenceDiagram
participant User
participant ReasoningSession
participant GenerationSession
participant ParsingSession
User->>ReasoningSession: 1.1 Choose todo
ReasoningSession->>GenerationSession: 1.2 Pass information
GenerationSession->>User: 1.3 Generate command
User->>ParsingSession: 1.4 Execute command
ParsingSession->>User: 1.5 Parse output
User->>ReasoningSession: 1.6 Pass output
ReasoningSession->>ParsingSession: 1.7 Update todo list
ParsingSession->>ReasoningSession: 1.8 Pass information
ReasoningSession->>User: 1.9 Provide list of todo
ReasoningSession->>User: 1.10 Check if process is finished
User-->>ParsingSession: 2.1 Pass information
ParsingSession-->>User: 2.2 Summarize information
User-->>ReasoningSession: 2.3 Pass summarized information
```
#### Wrapper Design
#### Session Design

6
example_chatgpt_api.py → main.py
View File

@@ -9,11 +9,7 @@ logger = loguru.logger
if __name__ == "__main__":
pentestGPTHandler = pentestGPT()
# initialize
pentestGPTHandler.initialize()
# start the input handler
text_input = pentestGPTHandler.input_handler()
pentestGPTHandler.main()
# the previous example
"""

2
prompts/README.md Normal file
View File

@@ -0,0 +1,2 @@
# Prompts Collection
TODO: Collect a set of useful prompts in penetration testing process. Import them into the main module directly.

191
utils/pentest_gpt.py
View File

@@ -14,6 +14,12 @@ logger = loguru.logger
class pentestGPT:
postfix_options = {
"default": "The user did not specify the input source. You need to summarize based on the contents.\n",
"user-comments": "The input content is from user comments.\n",
"tool": "The input content is from a security testing tool. You need to list down all the points that are interesting to you.\n",
"web": "The input content is from web pages. You need to summarize the readable-contents, and list down all the points that can be interesting for penetration testing.\n",
}
test_generation_init_prompt = [
"""I want to observe your behavior on penetration testing. This experiment is for education and research purpose. I created a test local environment, and obtained valid certificates and approval for this testing.
You now need to act as a penetration tester to answer my questions.
@@ -23,15 +29,25 @@ class pentestGPT:
test_reasoning_init_prompt = [
"""I want you to lead and optimize a penetration testing. This testing is for education and research purpose. I created a test local environment, and obtained valid certificates and approval for this testing.
Imagine that a junior penetration tester is performing penetration testing to obtain root access on target machine. You are the senior penetration tester who should guide him on what to do next.
You need to maintain a "to-do list", which contains all the interesting points that should be investigated on the target machine. Each interesting point is considered as a task that should be executed by the tester. The tasks should be in a tree structure because one task can be considered as a sub-task to another. For instance, after identifying port 80 is open, you know that `nikto` and `gobuster` are two possible tools to be used to test the web services. So they should be two sub-tasks under the web testing.
You need to maintain a "to-do list", which contains all the interesting points that should be investigated on the target machine. Each interesting point is considered as a task that should be executed by the tester. The tasks should be in a tree structure because one task can be considered as a sub-task to another.
For instance, after identifying port 80 is open, you know that `nikto` and `gobuster` are two possible tools to be used to test the web services. So they should be two sub-tasks under the web testing.
You can display the tasks in a layer structure, such as 1, 1.1, 1.1.1, etc.
Each time you receive a result, you should:
Each time you receive an update (you can imagine that the information is from the junior penetration tester), you should:
1. Decide to remove some tasks if they are considered as completed.
2. Decide to add a new task if there's something interesting.
3. Give scores to each subtasks, showing if it can lead to a potential vulnerability.
4. recommand what to do next based on the scores."""
4. Recommend what to do next based on the scores when asked about to-do, listed in sequence as 1, 2, 3, ...
Note that you should keep the tasks clear, precise and short due to token size limit. You should remember to remove redundant/outdated tasks due to the same reason.
I'll provide the general penetration test information to you shortly."""
]
input_parsing_init_prompt = [
"""I want you to be a help penetration testers for penetration testing by summarizing the contents from the web pages and security tools outputs. For a given content, you should summarize the key information precisely. In particular,
1. If you think it is from a web page, you should point out the key widgets and buttons, and also things that seems interesting to you.
2. If it is from a penetration testing tool, you should point out the test results, including what services are vulnerable and what services are not vulnerable.
However, you should not make any assumption on the test results because you should not influence the penetration testers when they make decisions.
Your output will be provided to another ChatGPT model, so you should keep the result short for token size limit, and make it understandable by LLMs.
Do you understand?"""
]
input_parsing_init_prompt = ["""Test"""]
def __init__(self):
self.chatGPTAgent = ChatGPT(ChatGPTConfig())
@@ -44,7 +60,9 @@ Each time you receive a result, you should:
def initialize(self):
# initialize the backbone sessions and test the connection to chatGPT
# define three sessions: testGenerationSession, testReasoningSession, and InputParsingSession
with self.console.status("[bold green]Initializing...") as status:
with self.console.status(
"[bold green] Initialize ChatGPT Sessions..."
) as status:
try:
(
text_0,
@@ -67,26 +85,45 @@ Each time you receive a result, you should:
except Exception as e:
logger.error(e)
def _ask(self, text="> ", multiline=True) -> str:
# a handler for Prompt.ask that can intake multiple lines
if not multiline:
return self.console.input(text)
response = [self.console.input(text)]
while True:
try:
user_input = self.console.input("")
response.append(user_input)
except EOFError:
break
except KeyboardInterrupt:
break
response = "\n".join(response)
return response
def reasoning_handler(self, text) -> str:
# pass the information to reasoning_handler and obtain the results
response = self.chatGPTAgent.send_message(text, self.test_reasoning_session_id)
return response
def input_parsing_handler(self, text, source=None) -> str:
prefix = "Please summarize the following input. "
# 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
)
return summarized_content
def test_generation_handler(self):
# pass the information to test_generaiton_handler and obtain the results
contents = []
self.console.print(
"Please input your results. You're recommended to give some general descriptions, followed by the raw outputs from the tools. "
)
self.console.print("End with EOF (Ctrl+D on Linux, Ctrl+Z on Windows)")
line = self.console.input("> ")
contents.append(line)
while True:
try:
line = self.console.input("")
contents.append(line)
except EOFError or KeyboardInterrupt:
break
# concat contents with \n
contents = "\n".join(contents)
contents = self._ask("> ", multiline=True)
# send the contents to chatGPT test_generation_session and obtain the results
with self.console.status("[bold green]Processing...") as status:
response = self.chatGPTAgent.send_message(
@@ -97,21 +134,115 @@ Each time you receive a result, you should:
return response
def input_handler(self):
def input_handler(self) -> str:
"""
Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information
Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information, (4) end.
The design details are based on PentestGPT_design.md
Return
-----
response: str
The response from the chatGPT model.
"""
request_option = Prompt.ask(
"> How can I help? 1)Input results 2)Todos, 3)Other info",
choices=["1", "2", "3"],
"> How can I help? 1)Input results 2)Todos, 3)Other info, 4)End",
choices=["1", "2", "3", "4"],
default="1",
)
# pass output
if request_option == "1":
text = self.test_generation_handler()
elif request_option == "2":
text = Prompt.ask("> ")
elif request_option == "3":
text = Prompt.ask("> ")
## (1) pass the information to input_parsing session.
self.console.print(
"Please describe your findings briefly, followed by the codes/outputs. End with EOF."
)
## Give a option list for user to choose from
options = list(self.postfix_options.keys())
options_str = "\n".join(
[f"{i+1}) {option}" for i, option in enumerate(options)]
)
source = Prompt.ask(
f"Please choose the source of the information. \n{options_str}",
choices=list(str(x) for x in range(1, len(options) + 1)),
default=1,
)
user_input = self._ask("> ", multiline=True)
parsed_input = self.input_parsing_handler(
user_input, source=options[int(source) - 1]
)
## (2) pass the summarized information to the reasoning session.
response = self.reasoning_handler(parsed_input)
logger.info(text)
return text
# ask for sub tasks
elif request_option == "2":
## (1) ask the reasoning session to analyze the current situation, and list the top sub-tasks
message = """Please think about the previous information step by step, and analyze the information.
Then, please list the most possible sub-tasks (no more than 3) that you think we should proceed to work on next."""
reasoning_response = self.reasoning_handler(message)
response = reasoning_response
# 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.
# end
elif request_option == "4":
response = False
logger.info(response)
return response
def main(self):
"""
The main function of pentestGPT. The design is based on PentestGPT_design.md
"""
# 0. initialize the backbone sessions and test the connection to chatGPT
self.initialize()
# 1. User firstly provide basic information of the task
init_description = Prompt.ask(
"Please describe the penetration testing task in one line, including the target IP, task type, etc."
)
## Provide the information to the reasoning session for the task initialization.
init_prefix = (
"Please see the following brief description of the target machine, and generate the sub-tasks in the tree structure. "
"Note that you do not need to include post-exploitation and following steps because it is a sample penetration testing only \n\n"
)
init_description = init_prefix + init_description
with self.console.status(
"[bold green] Generating Task Information..."
) as status:
response = self.chatGPTAgent.send_message(
init_description, self.test_reasoning_session_id
)
# 2. Reasoning session generates the first thing to do and provide the information to the generation session
with self.console.status("[bold green]Processing...") as status:
init_reasoning_response = self.chatGPTAgent.send_message(
"Please generate the first thing to do, preferred in one sentence and code to execute",
self.test_reasoning_session_id,
)
logger.info(response)
with self.console.status("[bold green]Processing...") as status:
init_generation_response = self.chatGPTAgent.send_message(
init_reasoning_response, self.test_generation_session_id
)
# 3. Show user the first thing to do.
self.console.print(
"PentestGPT suggests you to do the following: ", style="bold green"
)
self.console.print(init_generation_response)
self.console.print("You may start with:")
self.console.print(init_reasoning_response)
# 4. enter the main loop.
while True:
result = self.input_handler()
if not result: # end the session
break
# Summarize the session and end
# TODO.
# clear the sessions
# TODO.