"""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