mirror of
https://github.com/weyne85/PentestGPT.git
synced 2025-10-29 16:58:59 +00:00
118 lines
5.2 KiB
Python
118 lines
5.2 KiB
Python
# an automated penetration testing parser empowered by GPT
|
|
from config.chatgpt_config import ChatGPTConfig
|
|
from rich.spinner import Spinner
|
|
from utils.chatgpt import ChatGPT
|
|
from rich.prompt import Prompt
|
|
from rich.console import Console
|
|
|
|
|
|
import loguru
|
|
import time, os
|
|
|
|
|
|
logger = loguru.logger
|
|
|
|
|
|
class pentestGPT:
|
|
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 a result, 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."""
|
|
]
|
|
input_parsing_init_prompt = ["""Test"""]
|
|
|
|
def __init__(self):
|
|
self.chatGPTAgent = ChatGPT(ChatGPTConfig())
|
|
self.console = Console()
|
|
self.spinner = Spinner("line", "Processing")
|
|
self.test_generation_session_id = None
|
|
self.test_reasoning_session_id = None
|
|
self.input_parsing_session_id = None
|
|
|
|
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:
|
|
try:
|
|
(
|
|
text_0,
|
|
self.test_generation_session_id,
|
|
) = self.chatGPTAgent.send_new_message(
|
|
self.test_generation_init_prompt[0]
|
|
)
|
|
(
|
|
text_1,
|
|
self.test_reasoning_session_id,
|
|
) = self.chatGPTAgent.send_new_message(
|
|
self.test_reasoning_init_prompt[0]
|
|
)
|
|
(
|
|
text_2,
|
|
self.input_parsing_session_id,
|
|
) = self.chatGPTAgent.send_new_message(
|
|
self.input_parsing_init_prompt[0]
|
|
)
|
|
except Exception as e:
|
|
logger.error(e)
|
|
|
|
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)
|
|
|
|
# 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)
|
|
|
|
return response
|
|
|
|
def input_handler(self):
|
|
"""
|
|
Request for user's input to: (1) input test results, (2) ask for todos, (3) input other information
|
|
"""
|
|
request_option = Prompt.ask(
|
|
"> How can I help? 1)Input results 2)Todos, 3)Other info",
|
|
choices=["1", "2", "3"],
|
|
default="1",
|
|
)
|
|
if request_option == "1":
|
|
text = self.test_generation_handler()
|
|
elif request_option == "2":
|
|
text = Prompt.ask("> ")
|
|
elif request_option == "3":
|
|
text = Prompt.ask("> ")
|
|
|
|
logger.info(text)
|
|
return text
|