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: 🎸 update to v0.5

Browse Source 
Major update with local reasoning function and other doc updates
This commit is contained in:
Grey_D
2023-04-26 23:26:56 +08:00
parent 97d0fd1ad8
commit aca86402c7
7 changed files with 183 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
utils/chatgpt.py
View File

@@ -71,8 +71,7 @@ class ChatGPT:
# "cookie": f"cf_clearance={self.cf_clearance}; _puid={self._puid}; __Secure-next-auth.session-token={self.session_token}",
"cookie": self.config.cookie,
"user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36",
"accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7"
# 'Content-Type': 'text/event-stream; charset=utf-8',
"accept": "*/*",
}
)
self.headers["authorization"] = self.get_authorization()

130
utils/pentest_gpt.py
View File

@@ -6,7 +6,7 @@ from rich.console import Console
from prompts.prompt_class import PentestGPTPrompt
from utils.prompt_select import prompt_select, prompt_ask
from prompt_toolkit.formatted_text import HTML
from utils.task_handler import main_task_entry, mainTaskCompleter
from utils.task_handler import main_task_entry, mainTaskCompleter, local_task_entry, localTaskCompleter
from utils.web_parser import google_search, parse_web
import time
import datetime as dt
@@ -42,7 +42,7 @@ class pentestGPT:
"default": "The user did not specify the input source. You need to summarize based on the contents.\n",
}
def __init__(self, reasoning_model="gpt-4"):
def __init__(self, reasoning_model="text-davinci-002-render-sha"):
self.log_dir = "logs"
self.chatGPTAgent = ChatGPT(ChatGPTConfig())
self.chatGPT4Agent = ChatGPT(ChatGPTConfig(model=reasoning_model))
@@ -152,9 +152,95 @@ class pentestGPT:
self.log_conversation("generation", response)
return response
def local_input_handler(self) -> str:
"""
Request for user's input to handle the local task
"""
local_task_response = ""
self.chat_count += 1
local_request_option = local_task_entry()
self.log_conversation("user", local_request_option)
if local_request_option == "help":
print(localTaskCompleter().task_details)
elif local_request_option == "discuss":
## (1) Request for user multi-line input
self.console.print("Please share your findings and questions with PentestGPT.")
self.log_conversation(
"pentestGPT", "Please share your findings and questions with PentestGPT. (End with <shift + right-arrow>)"
)
user_input = prompt_ask(
"Your input: ", multiline=True
)
self.log_conversation("user", user_input)
## (2) pass the information to the reasoning session.
with self.console.status("[bold green] PentestGPT Thinking...") as status:
local_task_response = self.test_generation_handler(self.prompts.local_task_prefix + user_input)
## (3) print the results
self.console.print("PentestGPT:\n", style="bold green")
self.console.print(local_task_response + "\n", style="yellow")
self.log_conversation("pentestGPT", local_task_response)
elif local_request_option == "brainstorm":
## (1) Request for user multi-line input
self.console.print("Please share your concerns and questions with PentestGPT.")
self.log_conversation(
"pentestGPT", "Please share your concerns and questions with PentestGPT. End with <shift + right-arrow>)"
)
user_input = prompt_ask(
"Your input: ", multiline=True
)
self.log_conversation("user", user_input)
## (2) pass the information to the reasoning session.
with self.console.status("[bold green] PentestGPT Thinking...") as status:
local_task_response = self.test_generation_handler(self.prompts.local_task_brainstorm + user_input)
## (3) print the results
self.console.print("PentestGPT:\n", style="bold green")
self.console.print(local_task_response + "\n", style="yellow")
self.log_conversation("pentestGPT", local_task_response)
elif local_request_option == "google":
# get the users input
self.console.print(
"Please enter your search query. PentestGPT will summarize the info from google. (End with <shift + right-arrow>) ",
style="bold green",
)
self.log_conversation(
"pentestGPT",
"Please enter your search query. PentestGPT will summarize the info from google.",
)
user_input = prompt_ask(
"Your input: ", multiline=False
)
self.log_conversation("user", user_input)
with self.console.status("[bold green] PentestGPT Thinking...") as status:
# query the question
result: dict = google_search(user_input, 5) # 5 results by default
# summarize the results
# TODO
local_task_response = "Google search results:\n" + "still under development."
self.console.print(local_task_response + "\n", style="yellow")
self.log_conversation("pentestGPT", local_task_response)
return local_task_response
elif local_request_option == "continue":
self.console.print("Exit the local task and continue the main task.")
self.log_conversation("pentestGPT", "Exit the local task and continue the main task.")
local_task_response = "continue"
return local_task_response
def input_handler(self) -> str:
"""
Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information, (4) end.
Request for user's input to:
(1) input test results,
(2) ask for todos,
(3) input other information (discuss),
(4) google.
(4) end.
The design details are based on PentestGPT_design.md
Return
@@ -166,16 +252,6 @@ class pentestGPT:
request_option = main_task_entry()
self.log_conversation("user", request_option)
# request_option = prompt_select(
# title=f"({self.chat_count}) > Please select your options with cursor: ",
# values=[
# ("1", HTML('<style fg="cyan">Input test results</style>')),
# ("2", HTML('<style fg="cyan">Ask for todos</style>')),
# ("3", HTML('<style fg="cyan">Discuss with PentestGPT</style>')),
# ("4", HTML('<style fg="cyan">Exit</style>')),
# ],
# )
# pass output
if request_option == "help":
print(mainTaskCompleter().task_details)
@@ -222,7 +298,7 @@ class pentestGPT:
# generate more test details (beginner mode)
elif request_option == "more":
self.log_conversation("user", "more")
## (1) pass the reasoning results to the test_generation session.
## (1) check if reasoning session is initialized
if self.step_reasoning_response is None:
self.console.print(
"You have not initialized the task yet. Please perform the basic testing following `next` option.",
@@ -231,10 +307,20 @@ class pentestGPT:
response = "You have not initialized the task yet. Please perform the basic testing following `next` option."
self.log_conversation("pentestGPT", response)
return response
## (2) start local task generation.
### (2.1) ask the reasoning session to analyze the current situation, and explain the task
self.console.print("PentestGPT will generate more test details, and enter the sub-task generation mode. (Pressing Enter to continue)", style="bold green")
self.log_conversation("pentestGPT", "PentestGPT will generate more test details, and enter the sub-task generation mode.")
input()
### (2.2) pass the sub-tasks to the test generation session
with self.console.status("[bold green] PentestGPT Thinking...") as status:
generation_response = self.test_generation_handler(
self.step_reasoning_response
)
_local_init_response = self.test_generation_handler(
self.prompts.local_task_init
)
self.console.print(
"Below are the further details.",
@@ -244,6 +330,14 @@ class pentestGPT:
response = generation_response
self.log_conversation("pentestGPT", response)
### (2.3) local task handler
while True:
local_task_response = self.local_input_handler()
if local_task_response == "continue":
# break the local task handler
break
# ask for task list (to-do list)
elif request_option == "todo":
## log that user is asking for todo list
@@ -278,12 +372,12 @@ class pentestGPT:
# pass other information, such as questions or some observations.
elif request_option == "discuss":
## (1) Request for user multi-line input
self.console.print("Please share your thoughts/questions with PentestGPT.")
self.console.print("Please share your thoughts/questions with PentestGPT. (End with <shift + right-arrow>) ")
self.log_conversation(
"pentestGPT", "Please share your thoughts/questions with PentestGPT."
)
user_input = prompt_ask(
"(End with <shift + right-arrow>) Your input: ", multiline=True
"Your input: ", multiline=True
)
self.log_conversation("user", user_input)
## (2) pass the information to the reasoning session.
@@ -298,7 +392,7 @@ class pentestGPT:
elif request_option == "google":
# get the users input
self.console.print(
"Please enter your search query. PentestGPT will summarize the info from google.",
"Please enter your search query. PentestGPT will summarize the info from google. (End with <shift + right-arrow>) ",
style="bold green",
)
self.log_conversation(
@@ -306,7 +400,7 @@ class pentestGPT:
"Please enter your search query. PentestGPT will summarize the info from google.",
)
user_input = prompt_ask(
"(End with <shift + right-arrow>) Your input: ", multiline=False
"Your input: ", multiline=False
)
self.log_conversation("user", user_input)
with self.console.status("[bold green] PentestGPT Thinking...") as status:

49
utils/task_handler.py
View File

@@ -10,6 +10,43 @@ from prompt_toolkit.formatted_text import HTML
from prompt_toolkit.shortcuts import CompleteStyle, prompt
class localTaskCompleter(Completer):
tasks = [
"discuss", # discuss with pentestGPT on the local task
"brainstorm", # let pentestGPT brainstorm on the local task
"help", # show the help page (for this local task)
"google", # search on Google
"continue", # quit the local task (for this local task)
]
task_meta = {
"discuss": HTML("Discuss with <b>PentestGPT</b> about this local task."),
"brainstorm": HTML("Let <b>PentestGPT</b> brainstorm on the local task for all the possible solutions."),
"help": HTML("Show the help page for this local task."),
"google": HTML("Search on Google."),
"continue": HTML("Quit the local task and continue the previous testing."),
}
task_details = """
Below are the available tasks:
- discuss: Discuss with PentestGPT about this local task.
- brainstorm: Let PentestGPT brainstorm on the local task for all the possible solutions.
- help: Show the help page for this local task.
- google: Search on Google.
- quit: Quit the local task and continue the testing."""
def get_completions(self, document, complete_event):
word = document.get_word_before_cursor()
for task in self.tasks:
if task.startswith(word):
yield Completion(
task,
start_position=-len(word),
display=task,
display_meta=self.task_meta.get(task),
)
class mainTaskCompleter(Completer):
tasks = [
"next",
@@ -65,6 +102,18 @@ def main_task_entry(text="> "):
else:
return result
def local_task_entry(text="> "):
"""
Entry point for the task prompt. Auto-complete
"""
task_completer = localTaskCompleter()
while True:
result = prompt(text, completer=task_completer)
if result not in task_completer.tasks:
print("Invalid task, try again.")
else:
return result
if __name__ == "__main__":
main_task_entry()