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

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.