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: 🎸 separate prompts from script

Browse Source 
Now prompts under `prompts`. Also some minor bugs are fixed.

✅ Closes: #5
This commit is contained in:
Grey_D
2023-04-10 13:38:33 +08:00
parent 45ce3cccd6
commit 8d3a863db0
7 changed files with 207 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
utils/chatgpt.py
View File

@@ -14,6 +14,8 @@ import requests
from config.chatgpt_config import ChatGPTConfig
logger = loguru.logger
logger.remove()
logger.add(level = "WARNING", sink = "logs/chatgpt.log")
# A sample ChatGPTConfig class has the following structure. All fields can be obtained from the browser's cookie.
# In particular, cf_clearance、__Secure-next-auth.session-token、_puid are required.

129
utils/pentest_gpt.py
View File

@@ -4,53 +4,26 @@ from rich.spinner import Spinner
from utils.chatgpt import ChatGPT
from rich.prompt import Prompt
from rich.console import Console
from prompts.prompt_class import PentestGPTPrompt
import loguru
import time, os, textwrap
logger = loguru.logger
logger.add(sink="logs/pentest_gpt.log")
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",
"tool": "The input content is from a security testing tool. You need to list down all the points that are interesting to you; you should summarize it as if you are reporting to a senior penetration tester for further guidance.\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.
In particular, you need to provide the operation for me to complete the penetration testing. I'll describe the test scenarios for you.
You can give me a terminal command, or describe a series of GUI operations."""
]
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 can display the tasks in a layer structure, such as 1, 1.1, 1.1.1, etc.
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. 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?"""
]
def __init__(self):
self.chatGPTAgent = ChatGPT(ChatGPTConfig())
self.prompts = PentestGPTPrompt
self.console = Console()
self.spinner = Spinner("line", "Processing")
self.test_generation_session_id = None
@@ -68,20 +41,18 @@ Do you understand?"""
text_0,
self.test_generation_session_id,
) = self.chatGPTAgent.send_new_message(
self.test_generation_init_prompt[0]
self.prompts.generation_session_init
)
(
text_1,
self.test_reasoning_session_id,
) = self.chatGPTAgent.send_new_message(
self.test_reasoning_init_prompt[0]
self.prompts.reasoning_session_init
)
(
text_2,
self.input_parsing_session_id,
) = self.chatGPTAgent.send_new_message(
self.input_parsing_init_prompt[0]
)
) = self.chatGPTAgent.send_new_message(self.prompts.input_parsing_init)
except Exception as e:
logger.error(e)
@@ -116,8 +87,11 @@ Do you understand?"""
return response
def reasoning_handler(self, text) -> str:
# summarize the contents if necessary.
if len(text) > 8000:
text = self.input_parsing_handler(text)
# pass the information to reasoning_handler and obtain the results
response = self.chatGPTAgent.send_message(text, self.test_reasoning_session_id)
response = self.chatGPTAgent.send_message(self.prompts.process_results + text, self.test_reasoning_session_id)
return response
def input_parsing_handler(self, text, source=None) -> str:
@@ -135,27 +109,16 @@ Do you understand?"""
# (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"
word_limit = f"Please ensure that the input is less than {8000 / len(wrapped_inputs)} 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):
# pass the information to test_generaiton_handler and obtain the results
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)")
contents = self._ask("> ", multiline=True)
def test_generation_handler(self, text):
# 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(
contents, self.test_generation_session_id
)
# print the results
self.console.print(response)
response = self.chatGPTAgent.send_message(text, self.test_generation_session_id)
# print the results
return response
def input_handler(self) -> str:
@@ -194,30 +157,48 @@ Do you understand?"""
user_input, source=options[int(source) - 1]
)
## (2) pass the summarized information to the reasoning session.
response = self.reasoning_handler(parsed_input)
reasoning_response = self.reasoning_handler(parsed_input)
## (3) pass the reasoning results to the test_generation session.
generation_response = self.test_generation_handler(reasoning_response)
## (4) print the results
self.console.print("Based on the analysis, the following tasks are recommended:", style="bold green")
self.console.print(reasoning_response + '\n')
self.console.print("You can follow the instructions below to complete the tasks.", style="bold green")
self.console.print(generation_response + '\n')
response = generation_response
# 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)
reasoning_response = self.reasoning_handler(self.prompts.ask_todo)
## (2) pass the sub-tasks to the test_generation session.
message = self.prompts.todo_to_command + "\n" + reasoning_response
generation_response = self.test_generation_handler(message)
## (3) print the results
self.console.print("Based on the analysis, the following tasks are recommended:", style="bold green")
self.console.print(reasoning_response + '\n')
self.console.print("You can follow the instructions below to complete the tasks.", style="bold green")
self.console.print(generation_response + '\n')
response = reasoning_response
# pass other information, such as questions or some observations.
elif request_option == "3":
## (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)
## (2) pass the information to the reasoning session.
response = self.reasoning_handler(self.prompts.discussion + user_input)
## (3) print the results
self.console.print("PentestGPT:\n", style="bold green")
self.console.print(response + '\n', style="yellow")
# end
elif request_option == "4":
response = False
self.console.print("Thank you for using PentestGPT!", style="bold green")
logger.info(response)
return response
def main(self):
@@ -232,36 +213,24 @@ Do you understand?"""
"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
init_description = self.prompts.task_description + 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
)
_response = self.reasoning_handler(init_description)
# 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
first_todo = self.reasoning_handler(self.prompts.first_todo)
first_generation_response = self.test_generation_handler(
self.prompts.todo_to_command + first_todo
)
# 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)
self.console.print(first_todo)
self.console.print("You may start with:", style="bold green")
self.console.print(first_generation_response)
# 4. enter the main loop.
while True:

1
utils/web_parser.py Normal file
View File

@@ -0,0 +1 @@
# TODO: parse the web contents with bs4.