/pentestgpt/core/agent.py

"""Enhanced Claude Code agent with tracer integration for PentestGPT."""

import logging
import re
from pathlib import Path
from typing import Any, ClassVar

from claude_agent_sdk import (
    AssistantMessage,
    ClaudeAgentOptions,
    ClaudeSDKClient,
    CLIConnectionError,
    CLINotFoundError,
    ResultMessage,
    TextBlock,
    ToolUseBlock,
)

from pentestgpt.core.config import PentestGPTConfig
from pentestgpt.core.tracer import Tracer, get_global_tracer
from pentestgpt.prompts.pentesting import get_ctf_prompt

# Set up debug logging to file
# Try to write to workspace first, fallback to /tmp if permission denied
DEBUG_LOG_WORKSPACE = Path("/workspace/pentestgpt-debug.log")
DEBUG_LOG_FALLBACK = Path("/tmp/pentestgpt-debug.log")

handlers: list[logging.Handler] = [logging.StreamHandler()]  # Always log to stderr
try:
    # Try to create log file in workspace
    handlers.append(logging.FileHandler(DEBUG_LOG_WORKSPACE, mode="w"))
    DEBUG_LOG = DEBUG_LOG_WORKSPACE
except (PermissionError, OSError):
    # Fallback to /tmp if workspace is not writable
    handlers.append(logging.FileHandler(DEBUG_LOG_FALLBACK, mode="w"))
    DEBUG_LOG = DEBUG_LOG_FALLBACK
    print(
        f"Warning: Could not write to {DEBUG_LOG_WORKSPACE}, using {DEBUG_LOG_FALLBACK} instead",
        flush=True,
    )

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=handlers,
)
logger = logging.getLogger(__name__)


class PentestAgent:
    """Enhanced Claude Code agent with activity tracking and streaming updates."""

    # 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,
        tracer: Tracer | None = None,
    ) -> None:
        """
        Initialize the pentest agent.

        Args:
            config: Configuration for the agent
            tracer: Optional tracer instance (uses global if not provided)
        """
        self.config = config
        self.tracer = tracer or get_global_tracer()
        self.client: ClaudeSDKClient | None = None
        self.walkthrough_steps: list[str] = []
        self.flags_found: list[dict[str, str]] = []

    async def execute(self, task: str) -> dict[str, Any]:
        """
        Execute a penetration testing task.

        Args:
            task: The task/instructions for the agent

        Returns:
            Dictionary with execution results
        """
        logger.debug("=" * 80)
        logger.debug(f"EXECUTE START - Task: {task[:100]}...")
        self.tracer.track_agent_status("starting", "Initializing agent")

        try:
            # Setup Claude Code client
            logger.debug("Creating ClaudeAgentOptions...")
            options = ClaudeAgentOptions(
                cwd=str(self.config.working_directory),
                permission_mode=self.config.permission_mode,  # type: ignore[arg-type]
                system_prompt=self._build_system_prompt(),
                model=self.config.llm_model,
            )
            logger.debug(f"Options created: cwd={options.cwd}, model={options.model}")

            logger.debug("Creating ClaudeSDKClient...")
            self.client = ClaudeSDKClient(options=options)

            logger.debug("Attempting to connect to Claude Code CLI...")
            await self.client.connect()
            logger.debug("✓ Successfully connected to Claude Code CLI")

            self.tracer.track_agent_status("connected", "Connected to Claude Code")

            # Send the task
            logger.debug(f"Sending query to Claude Code: {task[:100]}...")
            await self.client.query(task)
            logger.debug("✓ Query sent, waiting for responses...")
            self.tracer.track_message(f"Task submitted: {task[:100]}...", message_type="info")

            # Process responses
            logger.debug("Starting message processing loop...")
            messages = []
            output_parts = []
            current_tool_id: int | None = None
            message_count = 0

            async for message in self.client.receive_response():
                message_count += 1
                logger.debug(f"← Received message #{message_count}: {type(message).__name__}")
                messages.append(message)

                if isinstance(message, AssistantMessage):
                    logger.debug(f"  AssistantMessage with {len(message.content)} blocks")
                    for block in message.content:
                        if isinstance(block, TextBlock):
                            logger.debug(f"  TextBlock: {block.text[:50]}...")
                            output_parts.append(block.text)

                            # Detect flags in the text
                            detected_flags = self._detect_flags(block.text)
                            if detected_flags:
                                for flag in detected_flags:
                                    flag_msg = f"🚩 FLAG DETECTED: {flag}"
                                    logger.info(flag_msg)
                                    self.tracer.track_message(flag_msg, message_type="success")
                                    self.flags_found.append(
                                        {"flag": flag, "context": block.text[:200]}
                                    )

                            self.tracer.track_message(block.text, message_type="info")

                            # Track as walkthrough step if it's substantial
                            if len(block.text.strip()) > 20:
                                self._add_walkthrough_step(block.text)

                        elif isinstance(block, ToolUseBlock):
                            logger.debug(f"  ToolUseBlock: {block.name}")
                            # Track tool execution start
                            current_tool_id = self.tracer.track_tool_start(
                                tool_name=block.name,
                                args=block.input,
                            )

                elif isinstance(message, ResultMessage):
                    logger.debug("  ResultMessage - Task completed")
                    # Mark last tool as completed if there was one
                    if current_tool_id is not None:
                        self.tracer.track_tool_complete(current_tool_id, status="completed")

                    # Track completion
                    cost = getattr(message, "total_cost_usd", 0)
                    logger.debug(f"  Total cost: ${cost:.4f}")
                    self.tracer.track_agent_status(
                        "completed", f"Task completed | Cost: ${cost:.4f}"
                    )

            logger.debug(f"Message loop ended. Processed {message_count} messages total.")

            result = {
                "success": True,
                "messages": messages,
                "output": "\n".join(output_parts),
                "cost_usd": getattr(messages[-1], "total_cost_usd", 0) if messages else 0,
                "walkthrough": self.walkthrough_steps,
                "flags_found": self.flags_found,
            }
            output_str = str(result["output"])
            logger.debug(
                f"✓ EXECUTE COMPLETE - Success: {result['success']}, Output length: {len(output_str)}"
            )
            logger.debug(f"  Flags found: {len(self.flags_found)}")
            logger.debug(f"  Walkthrough steps: {len(self.walkthrough_steps)}")
            return result

        except CLINotFoundError as e:
            logger.error(f"✗ CLINotFoundError: {e}")
            error_msg = (
                "Claude Code CLI not found. Install with: npm install -g @anthropic-ai/claude-code"
            )
            self.tracer.track_message(error_msg, message_type="error")
            return {"success": False, "error": error_msg}

        except CLIConnectionError as e:
            logger.error(f"✗ CLIConnectionError: {e}")
            error_msg = f"Failed to connect to Claude Code: {e!s}"
            self.tracer.track_message(error_msg, message_type="error")
            return {"success": False, "error": error_msg}

        except Exception as e:
            logger.error(f"✗ Unexpected Exception: {type(e).__name__}: {e}", exc_info=True)
            error_msg = f"Unexpected error: {e!s}"
            self.tracer.track_message(error_msg, message_type="error")
            return {"success": False, "error": error_msg}

        finally:
            logger.debug("Cleanup: disconnecting client...")
            if self.client:
                await self.client.disconnect()
                logger.debug("✓ Client disconnected")
            self.tracer.track_agent_status("disconnected", "Agent stopped")
            logger.debug("=" * 80)

    def _build_system_prompt(self) -> str:
        """Build the CTF system prompt for the agent."""
        return get_ctf_prompt(self.config.custom_instruction)

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

        Args:
            text: Text to search for flags

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

    def _add_walkthrough_step(self, step: str) -> None:
        """
        Add a step to the walkthrough.

        Args:
            step: Description of the step taken
        """
        # Clean up the step text
        step = step.strip()
        if step and step not in self.walkthrough_steps:
            self.walkthrough_steps.append(step)
            logger.debug(f"Walkthrough step added: {step[:100]}...")


async def run_pentest(
    target: str,
    custom_instruction: str | None = None,
    model: str | None = None,
    working_dir: str | None = None,
    debug: bool = False,
    resume_session: str | None = None,
) -> dict[str, Any]:
    """
    Convenience function to run a CTF challenge or penetration test.

    Uses the new AgentController for lifecycle management and session persistence.

    Args:
        target: Target challenge/machine to solve
        custom_instruction: Optional custom challenge context or instructions
        model: Optional model override
        working_dir: Optional working directory override
        debug: Enable debug mode with verbose console output
        resume_session: Optional session ID to resume

    Returns:
        Dictionary with challenge results including walkthrough and flags found
    """
    from pentestgpt.core.config import load_config
    from pentestgpt.core.controller import AgentController

    # Enable verbose console logging in debug mode
    if debug:
        logger.info("=" * 80)
        logger.info("DEBUG MODE ENABLED - CTF CHALLENGE SOLVER")
        logger.info(f"Debug log file: {DEBUG_LOG}")
        logger.info(f"Target: {target}")
        if custom_instruction:
            logger.info(f"Challenge context: {custom_instruction}")
        if model:
            logger.info(f"Model: {model}")
        if resume_session:
            logger.info(f"Resuming session: {resume_session}")
        logger.info("=" * 80)

    # Build config
    config_kwargs: dict[str, Any] = {"target": target}
    if custom_instruction:
        config_kwargs["custom_instruction"] = custom_instruction
    if model:
        config_kwargs["llm_model"] = model
    if working_dir:
        config_kwargs["working_directory"] = Path(working_dir)

    config = load_config(**config_kwargs)

    # Build task
    task = f"Solve this CTF challenge and capture the flag(s): {target}"
    if custom_instruction:
        task += f"\n\nChallenge context: {custom_instruction}"

    # Use AgentController for lifecycle management
    controller = AgentController(config)
    result = await controller.run(task, resume_session_id=resume_session)

    # Map result to legacy format for backward compatibility
    if result.get("success"):
        # Convert flags_found from list of strings to list of dicts if needed
        flags = result.get("flags_found", [])
        if flags and isinstance(flags[0], str):
            flags = [{"flag": f, "context": ""} for f in flags]

        result = {
            "success": True,
            "output": result.get("output", ""),
            "cost_usd": result.get("cost_usd", 0),
            "flags_found": flags,
            "walkthrough": [],  # Walkthrough tracking moved to session
            "session_id": result.get("session_id", ""),
        }

    if debug:
        logger.info("=" * 80)
        logger.info(f"CHALLENGE COMPLETE - Success: {result.get('success')}")
        if result.get("success"):
            logger.info(f"Flags found: {len(result.get('flags_found', []))}")
            logger.info(f"Cost: ${result.get('cost_usd', 0):.4f}")
            logger.info(f"Session: {result.get('session_id', 'N/A')}")
            for flag_data in result.get("flags_found", []):
                flag = (
                    flag_data.get("flag", flag_data) if isinstance(flag_data, dict) else flag_data
                )
                logger.info(f"  🚩 {flag}")
        else:
            logger.error(f"Error: {result.get('error', 'Unknown')}")
        logger.info("=" * 80)

    return result