# 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