run.py

"""
cyber-sim-mas — adaptive multi-agent cybersecurity simulation

Usage:
  python run.py --scenario lateral_movement --rounds 15 --verbose
  python run.py --scenario cloud_iam --seed 42
  python run.py --list-scenarios
"""

from __future__ import annotations

import argparse
import random
import sys

from rich import box
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

from cybersim.core.metrics import RoundSnapshot, SimulationMetrics
from cybersim.core.simulation import SimulationConfig, run_simulation
from cybersim.scenarios import REGISTRY

console = Console()


# ------------------------------------------------------------------
# CLI argument parsing
# ------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="run.py",
        description="Adaptive multi-agent cybersecurity simulation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  python run.py --scenario lateral_movement --rounds 15 --verbose
  python run.py --scenario cloud_iam --rounds 20 --seed 42
  python run.py --scenario supply_chain --rounds 12 --verbose --seed 7
  python run.py --list-scenarios
        """,
    )
    parser.add_argument(
        "--scenario",
        choices=list(REGISTRY.keys()),
        default="lateral_movement",
        metavar="NAME",
        help="scenario to run (default: lateral_movement)",
    )
    parser.add_argument(
        "--rounds",
        type=int,
        default=15,
        metavar="N",
        help="number of simulation rounds (default: 15)",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="print round-by-round activity",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        metavar="N",
        help="random seed for reproducible runs (omit for random)",
    )
    parser.add_argument(
        "--list-scenarios",
        action="store_true",
        help="list available scenarios and exit",
    )
    return parser


# ------------------------------------------------------------------
# Display helpers
# ------------------------------------------------------------------

def print_scenario_list() -> None:
    table = Table(title="Available Scenarios", box=box.SIMPLE_HEAD)
    table.add_column("Name", style="bold cyan")
    table.add_column("Description")
    for name, (_, description) in REGISTRY.items():
        table.add_row(name, description)
    console.print(table)


def print_round(
    round_num: int,
    snapshot: RoundSnapshot,
    attack_results: list,
    defender_actions: list,
) -> None:
    console.print(f"\n[bold]Round {round_num}[/bold]")

    # Attack activity
    for result in attack_results:
        icon = "[red]✗[/red]" if not result.success else "[green]✓[/green]"
        status = "success" if result.success else "failed"
        console.print(
            f"  {icon} [yellow]{result.technique.id}[/yellow] "
            f"{result.technique.name} → {result.target_node_id} ({status})"
        )

    # Defender actions
    for action in defender_actions:
        phase_color = {
            "investigate": "blue",
            "contain": "magenta",
            "remediate": "green",
        }.get(action.phase.value, "white")
        console.print(
            f"  [bold {phase_color}][{action.phase.value.upper()}][/bold {phase_color}] "
            f"{action.node_id} — {action.detail}"
        )

    # State summary
    if snapshot.compromised:
        console.print(
            f"  [dim]compromised: {', '.join(snapshot.compromised)}[/dim]"
        )


def print_summary(metrics: SimulationMetrics, scenario_name: str) -> None:
    data = metrics.to_dict()

    table = Table(title=f"Simulation Results — {scenario_name}", box=box.ROUNDED)
    table.add_column("Metric", style="bold")
    table.add_column("Value", justify="right")

    mttd = f"{int(data['mttd_rounds'])} rounds" if data["mttd_rounds"] is not None else "—"
    mttr = f"{int(data['mttr_rounds'])} rounds" if data["mttr_rounds"] is not None else "—"

    table.add_row("Total rounds", str(data["total_rounds"]))
    table.add_row("MTTD (mean time to detect)", mttd)
    table.add_row("MTTR (mean time to respond)", mttr)
    table.add_row(
        "Peak blast radius",
        f"{data['peak_blast_radius_pct']}%",
        style="red" if data["peak_blast_radius_pct"] > 50 else "",
    )
    table.add_row("Nodes compromised (final)", str(data["final_compromised"]))
    table.add_row("Attack success rate", f"{data['attack_success_rate_pct']}%")
    table.add_row(
        "Defender efficiency",
        f"{data['defender_efficiency_pct']}%",
        style="green" if data["defender_efficiency_pct"] > 60 else "red",
    )

    console.print()
    console.print(table)


# ------------------------------------------------------------------
# Entrypoint
# ------------------------------------------------------------------

def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.list_scenarios:
        print_scenario_list()
        sys.exit(0)

    if args.seed is not None:
        random.seed(args.seed)

    build_fn, description = REGISTRY[args.scenario]
    network, attacker, defender = build_fn()

    seed_label = f"seed={args.seed}" if args.seed is not None else "seed=random"
    console.print(
        Panel(
            f"[bold]{args.scenario.replace('_', ' ').title()}[/bold]\n"
            f"[dim]{description}[/dim]\n\n"
            f"Nodes: {len(network.nodes)}  ·  Rounds: {args.rounds}  ·  "
            f"Detection threshold: {defender.detection_threshold}  ·  "
            f"Techniques: {len(attacker.techniques)}  ·  {seed_label}",
            title="cyber-sim-mas",
            border_style="cyan",
        )
    )

    on_round = print_round if args.verbose else None

    config = SimulationConfig(
        scenario_name=args.scenario,
        rounds=args.rounds,
        verbose=args.verbose,
        on_round=on_round,
    )

    metrics = run_simulation(network, attacker, defender, config)
    print_summary(metrics, args.scenario)


if __name__ == "__main__":
    main()