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