mirror of
https://github.com/weyne85/PentestGPT.git
synced 2025-10-29 16:58:59 +00:00
42 lines
2.3 KiB
Python
42 lines
2.3 KiB
Python
# an automated penetration testing parser empowered by GPT
|
|
|
|
import loguru
|
|
from config.chatgpt_config import ChatGPTConfig
|
|
from utils.chatgpt import ChatGPT
|
|
|
|
logger = loguru.logger
|
|
|
|
|
|
class pentestGPT:
|
|
test_generation_init_prompt = ["""Test"""]
|
|
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())
|
|
# define three sessions: testGenerationSession, testReasoningSession, and InputParsingSession
|
|
text, test_generation_session_id = self.chatGPTAgent.send_new_message(self.test_generation_init_prompt[0])
|
|
text, test_reasoning_session_id = self.chatGPTAgent.send_new_message(self.test_reasoning_init_prompt[0])
|
|
text, input_parsing_session_id = self.chatGPTAgent.send_new_message(self.input_parsing_init_prompt[0])
|
|
|
|
|
|
def input_handler(self, text_input):
|
|
"""
|
|
Handle the user input from the terminal, and process it based on the input
|
|
------
|
|
input: text_input
|
|
output: text_output
|
|
"""
|
|
|
|
|