/pentestgpt/modules/core.py

# core.py

from rich.console import Console

from pentestgpt.modules.models import ModelInteraction


class CoreEngine:
    def __init__(self, state):
        self.state = state
        self.console = Console()
        self.print("Initializing the reasoning agent...", message_type="info")
        self.model = ModelInteraction(self, state)
        # Initialize other necessary components

    def run_iteration(self):
        # Gather the necessary states
        user_input = self.state.get_variable("user_input")
        task_description = self.state.get_variable("task_description")

        if not task_description:
            self.print(
                "Task description is not set. Please set it using 'set task_description [value]' command.",
                message_type="error",
            )
            return

        # Run the processing logic
        self.print("Processing the user input...", message_type="info")

        model_responses = self.model.process(user_input, task_description)
        self.state.log_conversation("pentestGPT", model_responses)
        self.print("PentestGPT Response:", message_type="info")

        for response in model_responses:
            self.print(response["message"], message_type=response["type"])

    def print(self, message, message_type=None, **kwargs):
        style_map = {
            "info": "bold green",
            "warning": "yellow",
            "error": "bold red",
            "heading": "bold magenta",
            "text": "",  # Default style
        }
        style = style_map.get(message_type, "")
        self.console.print(message, style=style)