# 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
        """