/pentestgpt/core/controller.py

"""Agent controller with lifecycle management, pause/resume, and session persistence."""

import asyncio
import re
from enum import Enum
from typing import Any, ClassVar

from pentestgpt.core.backend import (
    AgentBackend,
    AgentMessage,
    ClaudeCodeBackend,
    MessageType,
)
from pentestgpt.core.config import PentestGPTConfig
from pentestgpt.core.events import Event, EventBus, EventType
from pentestgpt.core.session import SessionStatus, SessionStore


class AgentState(Enum):
    """Simple 5-state model for agent lifecycle."""

    IDLE = "idle"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    ERROR = "error"


class AgentController:
    """
    Central orchestrator with lifecycle management.

    Features:
    - Framework-agnostic via AgentBackend
    - Pause/resume/stop control
    - Instruction injection
    - Session persistence
    """

    # Flag detection patterns
    FLAG_PATTERNS: ClassVar[list[str]] = [
        r"flag\{[^\}]+\}",  # flag{...}
        r"FLAG\{[^\}]+\}",  # FLAG{...}
        r"HTB\{[^\}]+\}",  # HTB{...}
        r"CTF\{[^\}]+\}",  # CTF{...}
        r"[A-Za-z0-9_]+\{[^\}]+\}",  # Generic CTF format
        r"\b[a-f0-9]{32}\b",  # 32-char hex (HTB user/root flags)
    ]

    def __init__(
        self,
        config: PentestGPTConfig,
        backend: AgentBackend | None = None,
        session_store: SessionStore | None = None,
        events: EventBus | None = None,
    ):
        """Initialize controller.

        Args:
            config: PentestGPT configuration
            backend: Optional custom backend (defaults to ClaudeCodeBackend)
            session_store: Optional custom session store
            events: Optional custom event bus
        """
        self.config = config
        self.backend = backend
        self.sessions = session_store or SessionStore()
        self.events = events or EventBus.get()

        # State management
        self._state = AgentState.IDLE
        self._pause_requested = False
        self._stop_requested = False
        self._resume_event = asyncio.Event()
        self._pending_instruction: str | None = None

        # Subscribe to user events
        self.events.subscribe(EventType.USER_COMMAND, self._on_user_command)
        self.events.subscribe(EventType.USER_INPUT, self._on_user_input)

    @property
    def state(self) -> AgentState:
        """Get current agent state."""
        return self._state

    def _set_state(self, state: AgentState, details: str = "") -> None:
        """Update state and emit event."""
        self._state = state
        self.events.emit_state(state.value, details)

    # === Control Methods (called from TUI) ===

    def pause(self) -> bool:
        """Request pause at next safe point.

        Returns:
            True if pause request was accepted
        """
        if self._state == AgentState.RUNNING:
            self._pause_requested = True
            return True
        return False

    def resume(self, instruction: str | None = None) -> bool:
        """Resume from paused state.

        Args:
            instruction: Optional instruction to inject on resume

        Returns:
            True if resume request was accepted
        """
        if self._state == AgentState.PAUSED:
            self._pending_instruction = instruction
            self._pause_requested = False
            self._resume_event.set()
            return True
        return False

    def stop(self) -> bool:
        """Request stop.

        Returns:
            True (stop is always accepted)
        """
        self._stop_requested = True
        self._resume_event.set()  # Unblock if paused
        return True

    def inject(self, instruction: str) -> bool:
        """Queue instruction for next pause point.

        Args:
            instruction: Instruction to inject

        Returns:
            True if instruction was queued
        """
        if self._state in (AgentState.RUNNING, AgentState.PAUSED):
            self._pending_instruction = instruction
            if self._state == AgentState.RUNNING:
                self._pause_requested = True
            return True
        return False

    # === Event Handlers ===

    def _on_user_command(self, event: Event) -> None:
        """Handle user command events."""
        cmd = event.data.get("command")
        if cmd == "pause":
            self.pause()
        elif cmd == "resume":
            self.resume()
        elif cmd == "stop":
            self.stop()

    def _on_user_input(self, event: Event) -> None:
        """Handle user input events."""
        text = event.data.get("text", "")
        if text:
            self.inject(text)

    # === Main Execution ===

    async def run(self, task: str, resume_session_id: str | None = None) -> dict[str, Any]:
        """Run agent with full lifecycle management.

        Args:
            task: Task description for the agent
            resume_session_id: Optional session ID to resume

        Returns:
            Result dictionary with success, output, flags, etc.
        """
        # Reset state
        self._pause_requested = False
        self._stop_requested = False
        self._resume_event.clear()

        # Create or resume session
        if resume_session_id:
            session = self.sessions.load(resume_session_id)
            if not session:
                return {
                    "success": False,
                    "error": f"Session {resume_session_id} not found",
                }
            # Update task if resuming
            if not task:
                task = session.task
        else:
            session = self.sessions.create(
                target=self.config.target,
                task=task,
                model=self.config.llm_model,
            )

        # Create backend if needed
        if self.backend is None:
            from pentestgpt.prompts.pentesting import get_ctf_prompt

            self.backend = ClaudeCodeBackend(
                working_directory=str(self.config.working_directory),
                system_prompt=get_ctf_prompt(self.config.custom_instruction),
                model=self.config.llm_model,
            )

        try:
            self._set_state(AgentState.RUNNING, "Connecting...")

            # Connect (or resume)
            if resume_session_id and self.backend.supports_resume:
                backend_session = session.backend_session_id or resume_session_id
                await self.backend.resume(backend_session)
                self.events.emit_message(f"Resumed session {resume_session_id}", "info")
            else:
                await self.backend.connect()

            # Store backend session ID if available
            if self.backend.session_id:
                self.sessions.set_backend_session_id(self.backend.session_id)

            # Send initial query
            await self.backend.query(task)
            self.sessions.update_status(SessionStatus.RUNNING)

            # Process messages with pause/stop handling
            output_parts: list[str] = []
            flags_found: list[str] = []

            async for msg in self.backend.receive_messages():
                # Check stop request
                if self._stop_requested:
                    self._set_state(AgentState.IDLE, "Stopped by user")
                    self.sessions.update_status(SessionStatus.PAUSED)
                    break

                # Check pause request (between messages = safe point)
                if self._pause_requested:
                    self._pause_requested = False
                    self._set_state(AgentState.PAUSED, "Paused - waiting for input")
                    self.sessions.update_status(SessionStatus.PAUSED)

                    # Wait for resume
                    await self._resume_event.wait()
                    self._resume_event.clear()

                    if self._stop_requested:
                        break

                    # Resume with pending instruction
                    self._set_state(AgentState.RUNNING, "Resumed")
                    self.sessions.update_status(SessionStatus.RUNNING)

                    if self._pending_instruction:
                        self.sessions.add_instruction(self._pending_instruction)
                        self.events.emit_message(
                            f"Injecting: {self._pending_instruction[:50]}...", "info"
                        )
                        await self.backend.query(self._pending_instruction)
                        self._pending_instruction = None

                # Process message by type
                await self._process_message(msg, output_parts, flags_found)

            # Completed successfully
            if not self._stop_requested:
                self._set_state(AgentState.COMPLETED)
                self.sessions.update_status(SessionStatus.COMPLETED)

            return {
                "success": True,
                "output": "\n".join(output_parts),
                "flags_found": flags_found,
                "session_id": session.session_id,
                "cost_usd": session.total_cost_usd,
            }

        except Exception as e:
            self._set_state(AgentState.ERROR, str(e))
            self.sessions.set_error(str(e))
            self.sessions.update_status(SessionStatus.ERROR)
            return {"success": False, "error": str(e)}

        finally:
            if self.backend:
                await self.backend.disconnect()

    async def _process_message(
        self, msg: AgentMessage, output_parts: list[str], flags_found: list[str]
    ) -> None:
        """Process a single agent message.

        Args:
            msg: Message to process
            output_parts: List to append text output to
            flags_found: List to append found flags to
        """
        if msg.type == MessageType.TEXT:
            output_parts.append(msg.content)
            self.events.emit_message(msg.content)

            # Detect flags
            detected = self._detect_flags(msg.content)
            for flag in detected:
                if flag not in flags_found:
                    flags_found.append(flag)
                    self.sessions.add_flag(flag, msg.content[:200])
                    self.events.emit_flag(flag, msg.content[:200])

        elif msg.type == MessageType.TOOL_START:
            self.events.emit_tool(
                status="start",
                name=msg.tool_name or "unknown",
                args=msg.tool_args,
            )

        elif msg.type == MessageType.TOOL_RESULT:
            self.events.emit_tool(
                status="complete",
                name=msg.tool_name or "unknown",
                result=msg.content,
            )

        elif msg.type == MessageType.RESULT:
            cost = msg.metadata.get("cost_usd", 0)
            if cost > 0:
                self.sessions.add_cost(cost)

    def _detect_flags(self, text: str) -> list[str]:
        """Detect potential flags in text.

        Args:
            text: Text to search for flags

        Returns:
            List of detected flag strings
        """
        flags = []
        for pattern in self.FLAG_PATTERNS:
            for match in re.finditer(pattern, text, re.IGNORECASE):
                flag = match.group(0)
                if flag not in flags:
                    flags.append(flag)
        return flags