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

fix: 🐛 login issue

Browse Source 
Add configurable user agent

✅ Closes: #26, #20, #13
This commit is contained in:
Grey_D
2023-04-27 11:00:45 +08:00
parent b70c5419e3
commit 5e957b5b1e
8 changed files with 64 additions and 484 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
utils/chatgpt.py
View File

@@ -68,11 +68,8 @@ class ChatGPT:
if not "cookie" in vars(self.config):
raise Exception("Please update cookie in config/chatgpt_config.py")
self.conversation_dict: Dict[str, Conversation] = {}
self.headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; rv:102.0) Gecko/20100101 Firefox/102.0',
'Accept': '*/*',
"Cookie": self.config.cookie
}
self.headers = {"Accept": "*/*", "Cookie": self.config.cookie}
self.headers["User-Agent"] = self.config.userAgent
self.headers["authorization"] = self.get_authorization()
def get_authorization(self):

73
utils/pentest_gpt.py
View File

@@ -6,7 +6,12 @@ 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, local_task_entry, localTaskCompleter
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
@@ -166,17 +171,20 @@ class pentestGPT:
elif local_request_option == "discuss":
## (1) Request for user multi-line input
self.console.print("Please share your findings and questions with PentestGPT.")
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
"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)
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")
@@ -184,23 +192,25 @@ class pentestGPT:
elif local_request_option == "brainstorm":
## (1) Request for user multi-line input
self.console.print("Please share your concerns and questions with PentestGPT.")
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
"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)
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(
@@ -211,28 +221,29 @@ class pentestGPT:
"pentestGPT",
"Please enter your search query. PentestGPT will summarize the info from google.",
)
user_input = prompt_ask(
"Your input: ", multiline=False
)
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."
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.")
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:
@@ -309,8 +320,14 @@ class pentestGPT:
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.")
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
@@ -372,13 +389,13 @@ 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. (End with <shift + right-arrow>) ")
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(
"Your input: ", multiline=True
)
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:
@@ -399,9 +416,7 @@ class pentestGPT:
"pentestGPT",
"Please enter your search query. PentestGPT will summarize the info from google.",
)
user_input = prompt_ask(
"Your input: ", multiline=False
)
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

5
utils/task_handler.py
View File

@@ -21,7 +21,9 @@ class localTaskCompleter(Completer):
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."),
"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."),
@@ -102,6 +104,7 @@ def main_task_entry(text="> "):
else:
return result
def local_task_entry(text="> "):
"""
Entry point for the task prompt. Auto-complete