From 7850f76bdb67678769cea1db64ccd8ece117bcd4 Mon Sep 17 00:00:00 2001 From: Grey_D Date: Wed, 5 Apr 2023 11:29:08 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20pentestGPT?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit add initial design doc for PentestGPT --- PentestGPT_design.md | 33 +++++++++++++++++++++++++++++++++ utils/pentest_gpt.py | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 PentestGPT_design.md create mode 100644 utils/pentest_gpt.py diff --git a/PentestGPT_design.md b/PentestGPT_design.md new file mode 100644 index 0000000..641a073 --- /dev/null +++ b/PentestGPT_design.md @@ -0,0 +1,33 @@ +# Design Documentation for PentestGPT +version 0.1, for web penetration testing only + +## General Design +PentestGPT provides a unified terminal input handler, and backed by three main components: +- A test generation module which generates the exact penetration testing commands or operations for the users to execute. +- A test reasoning module which conducts the reasoning of the test, guiding the penetration testers on what to do next. +- A parsing module which parses the output of the penetration tools and the contents on the webUI. + +## Function Design +The handler is the main entry point of the penetration testing tool. It allows pentesters to perform the following operations: +1. (initialize itself with some pre-designed prompts.) +2. Start a new penetration testing session by providing the target information. +3. Ask for todo-list, and acquire the next step to perform. +4. After completing the operation, pass the information to PentestGPT. + 1. Pass a tool output. + 2. Pass a webpage content. + 3. Pass a human description. + + +## System Design +### General Structure +1. Maintain three chat sessions in one class. Each session is for one component. +2. User can select to pass information to one section. In particular. + 1. todo: + 2. pass information: + +### Handler Design + +### Function Details + + + diff --git a/utils/pentest_gpt.py b/utils/pentest_gpt.py new file mode 100644 index 0000000..0f3cf18 --- /dev/null +++ b/utils/pentest_gpt.py @@ -0,0 +1,41 @@ +# 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 + """ + +