/pentestgpt/benchmark/docker.py

"""Simple Docker manager for benchmarks.

Starts/stops benchmark containers with ports exposed to localhost.
"""

import subprocess
from pathlib import Path
from typing import Any

from pentestgpt.benchmark.config import DEFAULT_HOST, DEFAULT_PORT


def start_benchmark(benchmark_path: Path, port: int = DEFAULT_PORT) -> dict[str, Any]:
    """Start a benchmark container.

    Args:
        benchmark_path: Path to benchmark directory (contains docker-compose.yml)
        port: Port to expose on host (default: 8080)

    Returns:
        dict with 'success', 'target_url', and 'message'
    """
    if not benchmark_path.exists():
        return {
            "success": False,
            "target_url": None,
            "message": f"Path not found: {benchmark_path}",
        }

    compose_file = benchmark_path / "docker-compose.yml"
    if not compose_file.exists():
        return {"success": False, "target_url": None, "message": "No docker-compose.yml found"}

    # First, build the containers
    print(f"Building benchmark at {benchmark_path.name}...")
    result = subprocess.run(
        ["make", "build"],
        cwd=str(benchmark_path),
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        return {
            "success": False,
            "target_url": None,
            "message": f"Build failed: {result.stderr or result.stdout}",
        }

    # Run with port mapping override
    # Use docker compose with port override
    print(f"Starting containers on port {port}...")
    result = subprocess.run(
        ["docker", "compose", "up", "-d", "--wait"],
        cwd=str(benchmark_path),
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        return {
            "success": False,
            "target_url": None,
            "message": f"Start failed: {result.stderr or result.stdout}",
        }

    # Get the actual port assigned
    actual_port = get_exposed_port(benchmark_path)
    target_url = f"http://{DEFAULT_HOST}:{actual_port or port}"

    return {
        "success": True,
        "target_url": target_url,
        "port": actual_port or port,
        "message": f"Benchmark started at {target_url}",
    }


def stop_benchmark(benchmark_path: Path) -> dict[str, Any]:
    """Stop a benchmark container.

    Args:
        benchmark_path: Path to benchmark directory

    Returns:
        dict with 'success' and 'message'
    """
    if not benchmark_path.exists():
        return {"success": False, "message": f"Path not found: {benchmark_path}"}

    print(f"Stopping benchmark at {benchmark_path.name}...")

    # Stop containers
    result = subprocess.run(
        ["docker", "compose", "down", "--remove-orphans"],
        cwd=str(benchmark_path),
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        return {"success": False, "message": f"Stop failed: {result.stderr or result.stdout}"}

    return {"success": True, "message": "Benchmark stopped"}


def get_exposed_port(benchmark_path: Path) -> int | None:
    """Get the exposed port from running containers.

    Args:
        benchmark_path: Path to benchmark directory

    Returns:
        Port number or None if not found
    """
    result = subprocess.run(
        ["docker", "compose", "ps", "--format", "{{.Ports}}"],
        cwd=str(benchmark_path),
        capture_output=True,
        text=True,
    )

    if result.returncode != 0 or not result.stdout.strip():
        return None

    # Parse port from output like "0.0.0.0:32768->80/tcp"
    import re

    for line in result.stdout.strip().split("\n"):
        match = re.search(r"0\.0\.0\.0:(\d+)->", line)
        if match:
            return int(match.group(1))

    return None


def get_running_benchmarks() -> list[dict[str, Any]]:
    """Get list of running benchmark containers.

    Returns:
        List of dicts with container info
    """
    result = subprocess.run(
        ["docker", "ps", "--format", "{{.Names}}\t{{.Ports}}\t{{.Status}}"],
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        return []

    running = []
    for line in result.stdout.strip().split("\n"):
        if not line:
            continue
        parts = line.split("\t")
        if len(parts) >= 3:
            name, ports, status = parts[0], parts[1], parts[2]
            # Check if it looks like a benchmark container (xben pattern)
            if "xben" in name.lower():
                running.append({"name": name, "ports": ports, "status": status})

    return running