/pentestgpt/modules/commands.py

# modules/commands.py

import shlex
from pentestgpt.utils.chat_utils.message import (
    AssistantMessage,
    SystemMessage,
    UserMessage,
)
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich.markdown import Markdown

class CommandParser:
    def __init__(self, state, core_engine):
        self.state = state
        self.core_engine = core_engine
        self.console = Console()
        self.commands = {
            "set": self.set_command,
            "show": self.show_command,
            "clear": self.clear_command,
            "run": self.run_command,
            "help": self.help_command,
            "exit": self.exit_command,
        }
        self.agent_names = ["reasoning", "parsing", "generation"]

    def parse(self, input_line):
        try:
            args = shlex.split(input_line)
            if not args:
                return
            command = args[0]
            cmd_args = args[1:]
            if command in self.commands:
                self.commands[command](cmd_args)
            else:
                self.console.print(
                    f"[bold red]Unknown command '{command}'. Type 'help' for a list of commands.[/bold red]"
                )
        except Exception as e:
            self.console.print(f"[bold red]Error parsing command: {str(e)}[/bold red]")

    def set_command(self, args):
        if len(args) < 2:
            self.console.print("[bold yellow]Usage: set [variable] [value][/bold yellow]")
            return
        key, value = args[0], " ".join(args[1:])
        self.state.set_variable(key, value)
        self.console.print(f"[green]Variable '{key}' set to '{value}'.[/green]")

    def show_command(self, args):
        if not args or args[0] == "all":
            variables = self.state.show_all()
            if variables:
                table = Table(title="Current Variables")
                table.add_column("Variable", style="cyan", no_wrap=True)
                table.add_column("Value", style="magenta")

                for key, value in variables.items():
                    table.add_row(key, value)

                self.console.print(table)
            else:
                self.console.print("[yellow]No variables set.[/yellow]")

            # Display available agents
            agent_list = ", ".join(self.agent_names)
            self.console.print("\n[bold]Available Agents:[/bold] [cyan]" + agent_list + "[/cyan]")
            self.console.print(
                "[bold]Use 'show [agent_name]' to display the agent's conversation history.[/bold]"
            )

        elif args[0] in self.agent_names:
            self.show_agent_history(args[0])
        else:
            key = args[0]
            value = self.state.get_variable(key)
            if value is not None:
                self.console.print(f"[bold]{key}:[/bold] {value}")
            else:
                self.console.print(f"[red]Variable or agent '{key}' is not recognized.[/red]")

    def show_agent_history(self, agent_name):
        agent = getattr(self.core_engine.model, f"{agent_name}_agent", None)
        if agent and agent.conversation_history:
            self.console.print(f"\n[bold underline]{agent_name.capitalize()} Agent Conversation History:[/bold underline]\n")
            for message in agent.conversation_history:
                role = message.role
                content = message.get_content()
                if role == "user":
                    panel = Panel(
                        content,
                        title="User",
                        title_align="left",
                        border_style="green",
                        expand=False,
                    )
                elif role == "assistant":
                    panel = Panel(
                        content,
                        title="Assistant",
                        title_align="left",
                        border_style="blue",
                        expand=False,
                    )
                elif role == "system":
                    panel = Panel(
                        content,
                        title="System",
                        title_align="left",
                        border_style="magenta",
                        expand=False,
                    )
                else:
                    panel = Panel(
                        content,
                        title=role.capitalize(),
                        title_align="left",
                        border_style="grey",
                        expand=False,
                    )
                self.console.print(panel)
        else:
            self.console.print(f"[yellow]No conversation history found for the {agent_name} agent.[/yellow]")

    def clear_command(self, args):
        if not args:
            self.console.print("[bold yellow]Usage: clear [variable][/bold yellow]")
            return
        key = args[0]
        self.state.clear_variable(key)
        self.console.print(f"[green]Variable '{key}' cleared.[/green]")

    def run_command(self, args):
        self.core_engine.run_iteration()

    def help_command(self, args):
        help_text = """
[bold]Available commands:[/bold]

- [cyan]set [variable] [value][/cyan]: Set a variable.
- [cyan]show [variable][/cyan]: Show the value of a variable.
  - Use [cyan]show all[/cyan] to display all variables and available agents.
  - Use [cyan]show [agent_name][/cyan] to display the agent's conversation history (e.g., 'show reasoning').
- [cyan]clear [variable][/cyan]: Clear a variable.
- [cyan]run[/cyan]: Execute the next iteration with current state.
- [cyan]help[/cyan]: Display this help message.
- [cyan]exit[/cyan]: Exit the application.
"""
        self.console.print(Markdown(help_text))

    def exit_command(self, args):
        self.console.print("[bold red]Exiting PentestGPT.[/bold red]")
        exit(0)