/pentestgpt/modules/state.py

# modules/state.py
# handler that maintains the internal states of all the variables in PentestGPT.


class State:
    def __init__(self):
        self.variables = {}
        self.history = {
            "user": [],
            "pentestGPT": [],
            "reasoning": [],
            "input_parsing": [],
            "generation": [],
            "exception": [],
        }

    def set_variable(self, key, value):
        self.variables[key] = value

    def get_variable(self, key):
        return self.variables.get(key, None)

    def show_all(self):
        return self.variables

    def clear_variable(self, key):
        if key in self.variables:
            del self.variables[key]

    def log_conversation(self, source, text):
        from time import time

        timestamp = time()
        if source not in self.history:
            source = "exception"
        self.history[source].append((timestamp, text))