mcghidra/src/ghydramcp/mixins/docker.py

"""Docker management mixin for GhydraMCP.

Provides tools for managing Ghidra Docker containers programmatically.
Allows the MCP server to automatically start containers when Ghidra isn't available.
"""

import asyncio
import os
import shutil
import subprocess
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastmcp import Context
from fastmcp.contrib.mcp_mixin import MCPMixin, mcp_tool

from ..config import get_config, get_docker_config


class DockerMixin(MCPMixin):
    """Docker container management for GhydraMCP.

    Provides tools to start, stop, and manage Ghidra containers
    with the GhydraMCP plugin pre-installed.
    """

    # Track running containers
    _containers: Dict[str, Dict[str, Any]] = {}

    def __init__(self):
        """Initialize Docker mixin."""
        self._check_docker_available()

    def _check_docker_available(self) -> bool:
        """Check if Docker is available on the system."""
        return shutil.which("docker") is not None

    def _run_docker_cmd(
        self, args: List[str], check: bool = True, capture: bool = True
    ) -> subprocess.CompletedProcess:
        """Run a docker command.

        Args:
            args: Command arguments (after 'docker')
            check: Raise exception on non-zero exit
            capture: Capture stdout/stderr

        Returns:
            CompletedProcess result
        """
        cmd = ["docker"] + args
        return subprocess.run(
            cmd,
            check=check,
            capture_output=capture,
            text=True,
        )

    def _run_compose_cmd(
        self,
        args: List[str],
        project_dir: Optional[Path] = None,
        check: bool = True,
        capture: bool = True,
    ) -> subprocess.CompletedProcess:
        """Run a docker compose command.

        Args:
            args: Command arguments (after 'docker compose')
            project_dir: Directory containing docker-compose.yml
            check: Raise exception on non-zero exit
            capture: Capture stdout/stderr

        Returns:
            CompletedProcess result
        """
        cmd = ["docker", "compose"]

        # Use project directory if specified
        if project_dir:
            cmd.extend(["-f", str(project_dir / "docker-compose.yml")])

        cmd.extend(args)

        env = os.environ.copy()
        if project_dir:
            env["COMPOSE_PROJECT_NAME"] = "ghydramcp"

        return subprocess.run(
            cmd,
            check=check,
            capture_output=capture,
            text=True,
            cwd=project_dir,
            env=env,
        )

    @mcp_tool(
        name="docker_status",
        description="Check Docker availability and running GhydraMCP containers",
    )
    async def docker_status(self, ctx: Optional[Context] = None) -> Dict[str, Any]:
        """Check Docker status and list running GhydraMCP containers.

        Returns:
            Status information including:
            - docker_available: Whether Docker is installed
            - docker_running: Whether Docker daemon is running
            - containers: List of GhydraMCP containers with their status
            - images: Available GhydraMCP images
        """
        result = {
            "docker_available": False,
            "docker_running": False,
            "containers": [],
            "images": [],
            "compose_available": False,
        }

        # Check if docker is installed
        if not self._check_docker_available():
            return result

        result["docker_available"] = True

        # Check if docker daemon is running
        try:
            self._run_docker_cmd(["info"], check=True)
            result["docker_running"] = True
        except (subprocess.CalledProcessError, FileNotFoundError):
            return result

        # Check for docker compose
        try:
            self._run_docker_cmd(["compose", "version"], check=True)
            result["compose_available"] = True
        except subprocess.CalledProcessError:
            pass

        # List GhydraMCP containers
        try:
            ps_result = self._run_docker_cmd(
                [
                    "ps",
                    "-a",
                    "--filter",
                    "label=org.opencontainers.image.title=ghydramcp",
                    "--format",
                    "{{.ID}}\t{{.Names}}\t{{.Status}}\t{{.Ports}}",
                ]
            )
            for line in ps_result.stdout.strip().split("\n"):
                if line:
                    parts = line.split("\t")
                    if len(parts) >= 3:
                        result["containers"].append(
                            {
                                "id": parts[0],
                                "name": parts[1],
                                "status": parts[2],
                                "ports": parts[3] if len(parts) > 3 else "",
                            }
                        )
        except subprocess.CalledProcessError:
            pass

        # Also check by name pattern
        try:
            ps_result = self._run_docker_cmd(
                [
                    "ps",
                    "-a",
                    "--filter",
                    "name=ghydramcp",
                    "--format",
                    "{{.ID}}\t{{.Names}}\t{{.Status}}\t{{.Ports}}",
                ]
            )
            existing_ids = {c["id"] for c in result["containers"]}
            for line in ps_result.stdout.strip().split("\n"):
                if line:
                    parts = line.split("\t")
                    if len(parts) >= 3 and parts[0] not in existing_ids:
                        result["containers"].append(
                            {
                                "id": parts[0],
                                "name": parts[1],
                                "status": parts[2],
                                "ports": parts[3] if len(parts) > 3 else "",
                            }
                        )
        except subprocess.CalledProcessError:
            pass

        # List GhydraMCP images
        try:
            images_result = self._run_docker_cmd(
                [
                    "images",
                    "--filter",
                    "reference=ghydramcp*",
                    "--format",
                    "{{.Repository}}:{{.Tag}}\t{{.Size}}\t{{.CreatedSince}}",
                ]
            )
            for line in images_result.stdout.strip().split("\n"):
                if line:
                    parts = line.split("\t")
                    if len(parts) >= 2:
                        result["images"].append(
                            {
                                "name": parts[0],
                                "size": parts[1],
                                "created": parts[2] if len(parts) > 2 else "",
                            }
                        )
        except subprocess.CalledProcessError:
            pass

        return result

    @mcp_tool(
        name="docker_start",
        description="Start a GhydraMCP Docker container to analyze a binary",
    )
    async def docker_start(
        self,
        binary_path: str,
        port: int = 8192,
        memory: str = "2G",
        name: Optional[str] = None,
        ctx: Optional[Context] = None,
    ) -> Dict[str, Any]:
        """Start a GhydraMCP Docker container for binary analysis.

        This creates a new Ghidra instance in Docker with the GhydraMCP
        plugin pre-installed. The binary will be imported and analyzed,
        then the HTTP API will be available on the specified port.

        Args:
            binary_path: Path to the binary file to analyze
            port: Port to expose the HTTP API (default: 8192)
            memory: Max JVM heap memory (default: 2G)
            name: Container name (auto-generated if not specified)

        Returns:
            Container info including ID, name, and API URL
        """
        if not self._check_docker_available():
            return {"error": "Docker is not available on this system"}

        # Verify binary exists
        binary_file = Path(binary_path).resolve()
        if not binary_file.exists():
            return {"error": f"Binary not found: {binary_path}"}

        # Generate container name if not specified
        if name is None:
            name = f"ghydramcp-{binary_file.stem}-{port}"

        # Clean up invalid characters in container name
        name = "".join(c if c.isalnum() or c in "-_" else "-" for c in name)

        try:
            # Check if container with this name already exists
            check_result = self._run_docker_cmd(
                ["ps", "-a", "-q", "-f", f"name={name}"], check=False
            )
            if check_result.stdout.strip():
                return {
                    "error": f"Container '{name}' already exists. Stop it first with docker_stop."
                }

            # Check if port is already in use
            port_check = self._run_docker_cmd(
                ["ps", "-q", "-f", f"publish={port}"], check=False
            )
            if port_check.stdout.strip():
                return {
                    "error": f"Port {port} is already in use by another container"
                }

            # Start the container
            run_result = self._run_docker_cmd(
                [
                    "run",
                    "-d",
                    "--name",
                    name,
                    "-p",
                    f"{port}:8192",
                    "-v",
                    f"{binary_file.parent}:/binaries:ro",
                    "-e",
                    f"GHYDRA_MAXMEM={memory}",
                    "ghydramcp:latest",
                    f"/binaries/{binary_file.name}",
                ]
            )

            container_id = run_result.stdout.strip()

            # Track the container
            self._containers[container_id] = {
                "name": name,
                "port": port,
                "binary": str(binary_file),
                "memory": memory,
            }

            return {
                "success": True,
                "container_id": container_id[:12],
                "name": name,
                "port": port,
                "api_url": f"http://localhost:{port}/",
                "binary": str(binary_file),
                "message": (
                    f"Container started. Analysis in progress. "
                    f"API will be available at http://localhost:{port}/ once analysis completes. "
                    f"Use docker_logs('{name}') to monitor progress."
                ),
            }

        except subprocess.CalledProcessError as e:
            return {"error": f"Failed to start container: {e.stderr or e.stdout}"}

    @mcp_tool(
        name="docker_stop",
        description="Stop a running GhydraMCP Docker container",
    )
    async def docker_stop(
        self, name_or_id: str, remove: bool = True, ctx: Optional[Context] = None
    ) -> Dict[str, Any]:
        """Stop a GhydraMCP Docker container.

        Args:
            name_or_id: Container name or ID
            remove: Also remove the container (default: True)

        Returns:
            Status message
        """
        if not self._check_docker_available():
            return {"error": "Docker is not available on this system"}

        try:
            # Stop the container
            self._run_docker_cmd(["stop", name_or_id])

            if remove:
                self._run_docker_cmd(["rm", name_or_id])
                # Remove from tracking
                self._containers = {
                    k: v
                    for k, v in self._containers.items()
                    if not (k.startswith(name_or_id) or v.get("name") == name_or_id)
                }
                return {"success": True, "message": f"Container '{name_or_id}' stopped and removed"}
            else:
                return {"success": True, "message": f"Container '{name_or_id}' stopped"}

        except subprocess.CalledProcessError as e:
            return {"error": f"Failed to stop container: {e.stderr or e.stdout}"}

    @mcp_tool(
        name="docker_logs",
        description="Get logs from a GhydraMCP Docker container",
    )
    async def docker_logs(
        self,
        name_or_id: str,
        tail: int = 100,
        follow: bool = False,
        ctx: Optional[Context] = None,
    ) -> Dict[str, Any]:
        """Get logs from a GhydraMCP container.

        Args:
            name_or_id: Container name or ID
            tail: Number of lines to show (default: 100)
            follow: Whether to follow log output (not recommended for MCP)

        Returns:
            Container logs
        """
        if not self._check_docker_available():
            return {"error": "Docker is not available on this system"}

        try:
            args = ["logs", "--tail", str(tail)]
            if follow:
                args.append("-f")
            args.append(name_or_id)

            result = self._run_docker_cmd(args)
            return {
                "success": True,
                "container": name_or_id,
                "logs": result.stdout + result.stderr,
            }

        except subprocess.CalledProcessError as e:
            return {"error": f"Failed to get logs: {e.stderr or e.stdout}"}

    @mcp_tool(
        name="docker_build",
        description="Build the GhydraMCP Docker image from source",
    )
    async def docker_build(
        self,
        tag: str = "latest",
        no_cache: bool = False,
        project_dir: Optional[str] = None,
        ctx: Optional[Context] = None,
    ) -> Dict[str, Any]:
        """Build the GhydraMCP Docker image.

        Args:
            tag: Image tag (default: 'latest')
            no_cache: Build without using cache
            project_dir: Path to GhydraMCP project (auto-detected if not specified)

        Returns:
            Build status
        """
        if not self._check_docker_available():
            return {"error": "Docker is not available on this system"}

        # Find project directory
        if project_dir:
            proj_path = Path(project_dir)
        else:
            # Try to find docker/Dockerfile relative to this file
            module_dir = Path(__file__).parent.parent.parent.parent
            if (module_dir / "docker" / "Dockerfile").exists():
                proj_path = module_dir
            else:
                return {
                    "error": "Could not find GhydraMCP project directory. Please specify project_dir."
                }

        dockerfile = proj_path / "docker" / "Dockerfile"
        if not dockerfile.exists():
            return {"error": f"Dockerfile not found at {dockerfile}"}

        try:
            args = [
                "build",
                "-t",
                f"ghydramcp:{tag}",
                "-f",
                str(dockerfile),
            ]
            if no_cache:
                args.append("--no-cache")
            args.append(str(proj_path))

            # Run build (this can take a while)
            result = self._run_docker_cmd(args, capture=True)

            return {
                "success": True,
                "image": f"ghydramcp:{tag}",
                "message": f"Successfully built ghydramcp:{tag}",
                "output": result.stdout[-2000:] if len(result.stdout) > 2000 else result.stdout,
            }

        except subprocess.CalledProcessError as e:
            return {"error": f"Build failed: {e.stderr or e.stdout}"}

    @mcp_tool(
        name="docker_health",
        description="Check if a GhydraMCP container's API is responding",
    )
    async def docker_health(
        self, port: int = 8192, timeout: float = 5.0, ctx: Optional[Context] = None
    ) -> Dict[str, Any]:
        """Check if a GhydraMCP container's API is healthy.

        Args:
            port: API port to check (default: 8192)
            timeout: Request timeout in seconds

        Returns:
            Health status and API info if available
        """
        import urllib.request
        import urllib.error
        import json

        url = f"http://localhost:{port}/"

        try:
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=timeout) as response:
                data = json.loads(response.read().decode())
                return {
                    "healthy": True,
                    "port": port,
                    "api_version": data.get("api_version"),
                    "program": data.get("program"),
                    "file": data.get("file"),
                }
        except urllib.error.URLError as e:
            return {
                "healthy": False,
                "port": port,
                "error": str(e.reason),
                "message": "Container may still be starting or analyzing binary",
            }
        except Exception as e:
            return {
                "healthy": False,
                "port": port,
                "error": str(e),
            }

    @mcp_tool(
        name="docker_wait",
        description="Wait for a GhydraMCP container to become healthy",
    )
    async def docker_wait(
        self,
        port: int = 8192,
        timeout: float = 300.0,
        interval: float = 5.0,
        ctx: Optional[Context] = None,
    ) -> Dict[str, Any]:
        """Wait for a GhydraMCP container to become healthy.

        Polls the API endpoint until it responds or timeout is reached.

        Args:
            port: API port to check (default: 8192)
            timeout: Maximum time to wait in seconds (default: 300)
            interval: Polling interval in seconds (default: 5)

        Returns:
            Health status once healthy, or error on timeout
        """
        import time

        start_time = time.time()
        last_error = None

        while (time.time() - start_time) < timeout:
            result = await self.docker_health(port=port, timeout=interval, ctx=ctx)
            if result.get("healthy"):
                result["waited_seconds"] = round(time.time() - start_time, 1)
                return result
            last_error = result.get("error")
            await asyncio.sleep(interval)

        return {
            "healthy": False,
            "port": port,
            "error": f"Timeout after {timeout}s waiting for container",
            "last_error": last_error,
        }

    @mcp_tool(
        name="docker_auto_start",
        description="Automatically start a GhydraMCP container if no Ghidra instance is available",
    )
    async def docker_auto_start(
        self,
        binary_path: str,
        port: int = 8192,
        wait: bool = True,
        timeout: float = 300.0,
        ctx: Optional[Context] = None,
    ) -> Dict[str, Any]:
        """Automatically start a Docker container if no Ghidra instance is available.

        This is the main entry point for automatic Docker management:
        1. Checks if a Ghidra instance is already running on the port
        2. If not, starts a new Docker container
        3. Optionally waits for the container to become healthy
        4. Returns connection info for the instance

        Args:
            binary_path: Path to the binary to analyze
            port: Port for the HTTP API (default: 8192)
            wait: Wait for container to be ready (default: True)
            timeout: Max wait time in seconds (default: 300)

        Returns:
            Instance connection info
        """
        # First, check if there's already a Ghidra instance on this port
        health = await self.docker_health(port=port, ctx=ctx)
        if health.get("healthy"):
            return {
                "source": "existing",
                "port": port,
                "api_url": f"http://localhost:{port}/",
                "program": health.get("program"),
                "message": "Using existing Ghidra instance",
            }

        # Check if Docker is available
        status = await self.docker_status(ctx=ctx)
        if not status.get("docker_running"):
            return {
                "error": "Docker is not available. Please install Docker or start Ghidra manually."
            }

        # Check if we have the image
        if not any("ghydramcp" in img.get("name", "") for img in status.get("images", [])):
            return {
                "error": (
                    "GhydraMCP Docker image not found. "
                    "Build it with docker_build() or 'make build' first."
                )
            }

        # Start a new container
        start_result = await self.docker_start(
            binary_path=binary_path, port=port, ctx=ctx
        )

        if not start_result.get("success"):
            return start_result

        if wait:
            # Wait for the container to become healthy
            wait_result = await self.docker_wait(port=port, timeout=timeout, ctx=ctx)
            if wait_result.get("healthy"):
                return {
                    "source": "docker",
                    "container_id": start_result.get("container_id"),
                    "container_name": start_result.get("name"),
                    "port": port,
                    "api_url": f"http://localhost:{port}/",
                    "program": wait_result.get("program"),
                    "waited_seconds": wait_result.get("waited_seconds"),
                    "message": f"Docker container ready after {wait_result.get('waited_seconds')}s",
                }
            else:
                return {
                    "warning": "Container started but not yet healthy",
                    "container_id": start_result.get("container_id"),
                    "port": port,
                    "last_error": wait_result.get("error"),
                    "message": "Container may still be analyzing. Check docker_logs() for progress.",
                }

        return {
            "source": "docker",
            "container_id": start_result.get("container_id"),
            "container_name": start_result.get("name"),
            "port": port,
            "api_url": f"http://localhost:{port}/",
            "message": "Container starting. Use docker_wait() or docker_health() to check status.",
        }