admin / AirportCyberSimulator

"""
Airport Cyber Resilience Simulator — Backend
=============================================
FastAPI application providing:
  * REST endpoints exposing the data-driven configuration (infrastructure,
    playbooks, compliance) loaded from /data.
  * A WebSocket-driven, server-authoritative simulation engine that steps
    through a selected playbook, emits ingress + timed spread events, applies
    mitigation trade-offs and streams live metrics to the kiosk UI.

Design goal: the engine is fully generic. Adding nodes, links, playbooks or
compliance mappings requires editing only the JSON files in /data — never this
module.
"""

from __future__ import annotations

import asyncio
import json
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
STATIC_DIR = BASE_DIR / "static"

# --------------------------------------------------------------------------- #
# Data loading (data-driven configuration layer)
# --------------------------------------------------------------------------- #


def load_json(name: str) -> Any:
    with (DATA_DIR / name).open("r", encoding="utf-8") as fh:
        return json.load(fh)


INFRASTRUCTURE = load_json("infrastructure.json")
PLAYBOOKS = load_json("playbooks.json")
COMPLIANCE = load_json("compliance.json")

PLAYBOOK_INDEX = {pb["playbook_id"]: pb for pb in PLAYBOOKS}
NODE_INDEX = {n["id"]: n for n in INFRASTRUCTURE["nodes"]}

# --------------------------------------------------------------------------- #
# Simulation engine
# --------------------------------------------------------------------------- #

# Economic / integrity tuning constants (kept here so behaviour is transparent).
COST_COMPROMISE = 250_000       # £ reputational + data loss per compromised node
COST_ISOLATE = 500_000          # £ operational loss to sever a node
COST_REDUCE = 200_000           # £ operational loss to throttle a node
COST_STOP_DEPLOY = 750_000      # £ to deploy an emergency patch team
SECURITY_HIT_PER_NODE = 15      # % security integrity lost per compromise
STOP_COUNTDOWN = 15             # seconds for the STOP ATTACK patch to complete
REDUCE_SLOWDOWN = 5.0           # remaining spread delay multiplier (80% slower)


class Simulation:
    """Server-authoritative state machine for a single client session."""

    def __init__(self, websocket: WebSocket) -> None:
        self.ws = websocket
        self.playbook: Optional[Dict[str, Any]] = None
        self.task: Optional[asyncio.Task] = None
        self.running = False

        # Live node state: id -> {"state": normal|ingress|compromised|isolated|
        #                          reduced|patched, "capacity": float}
        self.node_state: Dict[str, Dict[str, Any]] = {}

        # The ordered attack chain: list of dicts with source, target, delay,
        # vector, fired, cancelled. Index 0 sources from the ingress point.
        self.chain: List[Dict[str, Any]] = []

        # Active STOP-ATTACK patch efforts: source_node -> completion timestamp.
        self.patches: Dict[str, float] = {}

        self.metrics = {"security": 100.0, "capacity": 100.0, "financial": 0.0}
        self._reset_state()

    # -- helpers ----------------------------------------------------------- #

    def _reset_state(self) -> None:
        self.node_state = {
            nid: {"state": "normal", "capacity": 100.0} for nid in NODE_INDEX
        }
        self.chain = []
        self.patches = {}
        self.metrics = {"security": 100.0, "capacity": 100.0, "financial": 0.0}

    def _recompute_capacity(self) -> None:
        total_w = sum(n["ops_weight"] for n in INFRASTRUCTURE["nodes"])
        weighted = 0.0
        for n in INFRASTRUCTURE["nodes"]:
            weighted += self.node_state[n["id"]]["capacity"] * n["ops_weight"]
        self.metrics["capacity"] = round(weighted / total_w, 1)

    async def send(self, payload: Dict[str, Any]) -> None:
        try:
            await self.ws.send_text(json.dumps(payload))
        except (WebSocketDisconnect, RuntimeError):
            self.running = False

    async def _emit_metrics(self) -> None:
        self._recompute_capacity()
        self.metrics["security"] = round(max(0.0, self.metrics["security"]), 1)
        await self.send({"type": "metrics", "metrics": self.metrics})

    # -- lifecycle --------------------------------------------------------- #

    async def start(self, playbook_id: str) -> None:
        await self.stop_task()
        pb = PLAYBOOK_INDEX.get(playbook_id)
        if not pb:
            await self.send({"type": "error", "message": f"Unknown playbook {playbook_id}"})
            return

        self._reset_state()
        self.playbook = pb
        self.running = True

        # Build the attack chain. Each spread event's source is the previous
        # compromised node (ingress for the first hop).
        prev = pb["ingress_point"]
        for ev in pb["spread_events"]:
            self.chain.append(
                {
                    "source": prev,
                    "target": ev["target"],
                    "delay": float(ev["delay_seconds"]),
                    "vector": ev["vector"],
                    "fired": False,
                    "cancelled": False,
                }
            )
            prev = ev["target"]

        await self.send(
            {
                "type": "sim_start",
                "playbook": pb,
                "ingress_node": pb["ingress_point"],
            }
        )

        self.task = asyncio.create_task(self._run())

    async def _run(self) -> None:
        """Drive the timeline: ingress, then timed spread events."""
        pb = self.playbook
        assert pb is not None
        ingress = pb["ingress_point"]

        # 1. Show the method-of-ingress indicator over the point of ingress.
        await self.send(
            {
                "type": "ingress",
                "node": ingress,
                "method": pb["ingress_method"],
                "icon": pb["ingress_icon"],
            }
        )
        self.node_state[ingress]["state"] = "ingress"
        await asyncio.sleep(2.0)  # dramatic pause so the method is readable

        # 2. Ingress node becomes compromised.
        await self._compromise(ingress, vector=pb["ingress_method"], initial=True)

        # 3. Walk the timeline of spread events using absolute delays.
        start_t = time.monotonic()
        idx = 0
        while self.running and idx < len(self.chain):
            hop = self.chain[idx]
            # Wait until this hop's (possibly mutated) delay elapses.
            while self.running:
                if hop["cancelled"]:
                    break
                elapsed = time.monotonic() - start_t
                if elapsed >= hop["delay"]:
                    break
                await self._tick_patches()
                await asyncio.sleep(0.1)

            if not self.running:
                return
            if hop["cancelled"]:
                await self.send(
                    {"type": "spread_blocked", "source": hop["source"], "target": hop["target"]}
                )
                idx += 1
                continue

            # Fire the spread event if the source is still an active threat.
            src_state = self.node_state[hop["source"]]["state"]
            if src_state in ("isolated", "patched"):
                hop["cancelled"] = True
                await self.send(
                    {"type": "spread_blocked", "source": hop["source"], "target": hop["target"]}
                )
                idx += 1
                continue

            await self.send(
                {
                    "type": "spread",
                    "source": hop["source"],
                    "target": hop["target"],
                    "vector": hop["vector"],
                }
            )
            hop["fired"] = True
            await self._compromise(hop["target"], vector=hop["vector"])
            idx += 1

        if self.running:
            await self.send({"type": "sim_complete"})

    async def _tick_patches(self) -> None:
        """Resolve any STOP-ATTACK countdowns that have completed."""
        now = time.monotonic()
        done = [src for src, ts in self.patches.items() if now >= ts]
        for src in done:
            del self.patches[src]
            # Cancel every not-yet-fired hop sourced from the patched node.
            blocked = False
            for hop in self.chain:
                if hop["source"] == src and not hop["fired"] and not hop["cancelled"]:
                    hop["cancelled"] = True
                    blocked = True
            self.node_state[src]["state"] = "patched"
            self.metrics["security"] = min(100.0, self.metrics["security"] + 10)
            await self.send(
                {
                    "type": "patch_complete",
                    "node": src,
                    "success": blocked,
                }
            )
            await self._emit_metrics()

    async def _compromise(self, node: str, vector: str, initial: bool = False) -> None:
        st = self.node_state[node]
        if st["state"] in ("compromised", "isolated", "patched"):
            return
        st["state"] = "compromised"
        st["capacity"] = 70.0  # degraded but still limping along
        self.metrics["security"] -= SECURITY_HIT_PER_NODE
        self.metrics["financial"] += COST_COMPROMISE
        await self.send(
            {
                "type": "compromised",
                "node": node,
                "vector": vector,
                "initial": initial,
                "compliance": COMPLIANCE["mappings"].get(node),
            }
        )
        await self._emit_metrics()

    # -- mitigations ------------------------------------------------------- #

    async def mitigate(self, node: str, option: str) -> None:
        if node not in self.node_state:
            return
        option = option.upper()
        st = self.node_state[node]

        if option == "ISOLATE":
            st["state"] = "isolated"
            st["capacity"] = 0.0
            self.metrics["financial"] += COST_ISOLATE
            # Sever every hop leaving this node.
            for hop in self.chain:
                if hop["source"] == node and not hop["fired"]:
                    hop["cancelled"] = True
            await self.send(
                {"type": "mitigation", "node": node, "option": "ISOLATE",
                 "message": f"{NODE_INDEX[node]['name']} isolated — links severed, spread halted."}
            )

        elif option == "REDUCE":
            st["state"] = "reduced"
            st["capacity"] = 40.0
            self.metrics["financial"] += COST_REDUCE
            # Slow the remaining outbound hops by 80%.
            for hop in self.chain:
                if hop["source"] == node and not hop["fired"] and not hop["cancelled"]:
                    hop["delay"] *= REDUCE_SLOWDOWN
            await self.send(
                {"type": "mitigation", "node": node, "option": "REDUCE",
                 "message": f"{NODE_INDEX[node]['name']} traffic reduced — spread slowed by 80%."}
            )

        elif option == "STOP":
            self.metrics["financial"] += COST_STOP_DEPLOY
            self.patches[node] = time.monotonic() + STOP_COUNTDOWN
            await self.send(
                {"type": "patch_start", "node": node, "seconds": STOP_COUNTDOWN,
                 "message": f"Emergency patch deploying on {NODE_INDEX[node]['name']} — {STOP_COUNTDOWN}s."}
            )
        else:
            return

        await self._emit_metrics()

    async def stop_task(self) -> None:
        self.running = False
        if self.task and not self.task.done():
            self.task.cancel()
            try:
                await self.task
            except asyncio.CancelledError:
                pass
        self.task = None

    async def reset(self) -> None:
        await self.stop_task()
        self._reset_state()
        await self.send({"type": "reset"})
        await self._emit_metrics()


# --------------------------------------------------------------------------- #
# FastAPI wiring
# --------------------------------------------------------------------------- #

app = FastAPI(title="Airport Cyber Resilience Simulator")


@app.get("/api/infrastructure")
async def api_infrastructure() -> JSONResponse:
    return JSONResponse(INFRASTRUCTURE)


@app.get("/api/playbooks")
async def api_playbooks() -> JSONResponse:
    return JSONResponse(PLAYBOOKS)


@app.get("/api/compliance")
async def api_compliance() -> JSONResponse:
    return JSONResponse(COMPLIANCE)


@app.get("/api/bootstrap")
async def api_bootstrap() -> JSONResponse:
    """Single call the frontend uses to hydrate the whole kiosk."""
    return JSONResponse(
        {
            "infrastructure": INFRASTRUCTURE,
            "playbooks": PLAYBOOKS,
            "compliance": COMPLIANCE,
        }
    )


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket) -> None:
    await ws.accept()
    sim = Simulation(ws)
    await sim._emit_metrics()
    try:
        while True:
            raw = await ws.receive_text()
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                continue
            action = msg.get("action")
            if action == "start":
                await sim.start(msg.get("playbook_id", ""))
            elif action == "mitigate":
                await sim.mitigate(msg.get("node", ""), msg.get("option", ""))
            elif action == "reset":
                await sim.reset()
            elif action == "ping":
                await sim.send({"type": "pong"})
    except WebSocketDisconnect:
        await sim.stop_task()
    except Exception:
        await sim.stop_task()


# Serve the kiosk UI (mounted last so /api and /ws take precedence).
app.mount("/", StaticFiles(directory=str(STATIC_DIR), html=True), name="static")