/pentestgpt/benchmark/cli.py

"""Simple CLI for benchmark management.

Usage:
    pentestgpt-benchmark list [--tags TAG ...] [--levels N ...]
    pentestgpt-benchmark start BENCHMARK_ID
    pentestgpt-benchmark stop BENCHMARK_ID
    pentestgpt-benchmark status
"""

import argparse
import sys
from pathlib import Path

from pentestgpt.benchmark.config import DEFAULT_BENCHMARKS_DIR
from pentestgpt.benchmark.docker import (
    get_running_benchmarks,
    start_benchmark,
    stop_benchmark,
)
from pentestgpt.benchmark.registry import BenchmarkRegistry


def cmd_list(args: argparse.Namespace) -> int:
    """List available benchmarks."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    try:
        registry = BenchmarkRegistry(benchmarks_dir)
        registry.load()
    except FileNotFoundError as e:
        print(f"Error: {e}")
        return 1

    # Show tags if requested
    if args.show_tags:
        print("Available tags:")
        for tag in sorted(registry.get_all_tags()):
            count = len(registry.filter(tags=[tag]))
            print(f"  {tag}: {count} benchmarks")
        return 0

    # Filter and list
    benchmarks = registry.filter(tags=args.tags, levels=args.levels)

    print(f"{'ID':<15} {'Level':^5} {'Name':<40} {'Tags'}")
    print("-" * 80)
    for b in benchmarks:
        tags_str = ", ".join(b.tags[:3])
        if len(b.tags) > 3:
            tags_str += "..."
        print(f"{b.id:<15} {b.level:^5} {b.name[:40]:<40} {tags_str}")

    print(f"\nTotal: {len(benchmarks)} benchmarks")
    return 0


def cmd_start(args: argparse.Namespace) -> int:
    """Start a benchmark."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    registry = BenchmarkRegistry(benchmarks_dir)
    benchmark = registry.get(args.benchmark_id)

    if not benchmark:
        print(f"Error: Benchmark '{args.benchmark_id}' not found")
        return 1

    print(f"Starting benchmark: {benchmark.name}")
    print(f"Description: {benchmark.description}")
    print(f"Level: {benchmark.level}, Tags: {', '.join(benchmark.tags)}")
    print()

    result = start_benchmark(benchmark.path)

    if result["success"]:
        target_url = result["target_url"]
        # Extract port from URL for Docker target
        port = target_url.split(":")[-1]
        docker_url = f"http://host.docker.internal:{port}"

        print("\nBenchmark started successfully!")
        print(f"Target URL: {target_url}")
        print()
        print("Run PentestGPT against this target:")
        print(f"  Local:  pentestgpt --target {target_url}")
        print(f"  Docker: pentestgpt --target {docker_url}")
    else:
        print(f"\nFailed to start benchmark: {result['message']}")
        return 1

    return 0


def cmd_stop(args: argparse.Namespace) -> int:
    """Stop a benchmark."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    registry = BenchmarkRegistry(benchmarks_dir)
    benchmark = registry.get(args.benchmark_id)

    if not benchmark:
        print(f"Error: Benchmark '{args.benchmark_id}' not found")
        return 1

    result = stop_benchmark(benchmark.path)

    if result["success"]:
        print("Benchmark stopped successfully")
    else:
        print(f"Failed to stop benchmark: {result['message']}")
        return 1

    return 0


def cmd_status(args: argparse.Namespace) -> int:
    """Show running benchmarks."""
    running = get_running_benchmarks()

    if not running:
        print("No benchmark containers currently running")
        return 0

    print("Running benchmark containers:")
    print(f"{'Name':<40} {'Ports':<30} {'Status'}")
    print("-" * 80)
    for container in running:
        print(f"{container['name']:<40} {container['ports']:<30} {container['status']}")

    return 0


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        prog="pentestgpt-benchmark",
        description="Manage benchmark containers for PentestGPT testing",
    )
    parser.add_argument(
        "--benchmarks-dir",
        "-d",
        help="Path to benchmarks directory",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # List command
    list_parser = subparsers.add_parser("list", help="List available benchmarks")
    list_parser.add_argument("--tags", "-t", nargs="+", help="Filter by tags")
    list_parser.add_argument(
        "--levels", "-l", nargs="+", type=int, help="Filter by levels (1, 2, 3)"
    )
    list_parser.add_argument("--show-tags", action="store_true", help="Show all tags")

    # Start command
    start_parser = subparsers.add_parser("start", help="Start a benchmark")
    start_parser.add_argument("benchmark_id", help="Benchmark ID (e.g., XBEN-001-24)")

    # Stop command
    stop_parser = subparsers.add_parser("stop", help="Stop a benchmark")
    stop_parser.add_argument("benchmark_id", help="Benchmark ID")

    # Status command
    subparsers.add_parser("status", help="Show running benchmarks")

    args = parser.parse_args()

    if args.command == "list":
        sys.exit(cmd_list(args))
    elif args.command == "start":
        sys.exit(cmd_start(args))
    elif args.command == "stop":
        sys.exit(cmd_stop(args))
    elif args.command == "status":
        sys.exit(cmd_status(args))


if __name__ == "__main__":
    main()