whitehat-security/src/whitehat_security/plane.py

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any, Protocol

from .engagement import AuthorizationError, Engagement

PLANE_NAMESPACE = "whitehat"
STANDING_E2_TTL = 900
TERMINAL_IDS = frozenset({
    "WH-ENG-20260821-AUDIT-E2",
    "WH-ENG-20260821-TENANT-E2",
})
APPROVAL_CLASSES = {
    "fixture-e2": frozenset({"e2-differential"}),
    "live-e2": frozenset({"e2-differential"}),
    "fixture-e3": frozenset({"e3-rls"}),
    "e3": frozenset({"e3-rls"}),
    "fixture-capacity": frozenset({"p1-noisy-neighbour", "p2-noisy-neighbour"}),
    "capacity": frozenset({"p1-noisy-neighbour", "p2-noisy-neighbour"}),
}
FIXTURE_CLASSES = frozenset({"fixture-e2", "fixture-e3", "fixture-capacity"})


def repo_root() -> Path:
    return Path(__file__).resolve().parents[2]


def retired_ids(engagements_dir: str | Path | None = None) -> set[str]:
    retired = set(TERMINAL_IDS)
    root = Path(engagements_dir) if engagements_dir else repo_root() / "engagements"
    if not root.exists():
        return retired
    for path in root.glob("*.json"):
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError, ValueError):
            continue
        if data.get("status") == "cancelled" and data.get("engagement_id"):
            retired.add(data["engagement_id"])
    return retired


class KillSwitch:
    def __init__(self, path: str | Path | None = None) -> None:
        default = os.environ.get("WHITEHAT_KILL_SWITCH_PATH", str(repo_root() / "plane" / "KILL"))
        self.path = Path(path) if path is not None else Path(default)

    def engaged(self) -> bool:
        flag = os.environ.get("WHITEHAT_KILL_SWITCH", "").strip().lower()
        if flag in {"1", "true", "yes", "on"}:
            return True
        return self.path.exists()


class RateWatcher:
    def __init__(self, *, per_minute: int, max_concurrency: int, max_requests: int) -> None:
        if max_concurrency != 1:
            raise AuthorizationError("v0.1 permits exactly one in-flight API/database operation")
        if not 1 <= per_minute <= 60:
            raise AuthorizationError("API rate ceiling must be between 1 and 60 per minute")
        if max_requests < 1:
            raise AuthorizationError("request ceiling must be at least 1")
        self.per_minute = per_minute
        self.max_requests = max_requests
        self._times: list[float] = []
        self._in_flight = 0
        self._count = 0

    def acquire(self) -> None:
        if self._in_flight:
            raise AuthorizationError("concurrency ceiling reached")
        now = time.monotonic()
        self._times = [stamp for stamp in self._times if now - stamp < 60]
        if len(self._times) >= self.per_minute:
            raise AuthorizationError("rate ceiling reached")
        if self._count >= self.max_requests:
            raise AuthorizationError("request ceiling reached")
        self._times.append(now)
        self._count += 1
        self._in_flight = 1

    def release(self) -> None:
        self._in_flight = 0


@dataclass(frozen=True)
class IdentityHandle:
    role: str
    mount_path: str
    lease_id: str
    expires_at: str

    def __repr__(self) -> str:
        return (
            f"IdentityHandle(role={self.role!r}, mount_path={self.mount_path!r}, "
            f"lease_id={self.lease_id!r}, expires_at={self.expires_at!r})"
        )


class CredentialBroker(Protocol):
    def project(self, engagement: Engagement, registration: dict[str, Any]
                ) -> tuple[IdentityHandle, ...]:
        ...

    def revoke(self, lease_id: str) -> None:
        ...


class UnconnectedCustodyBroker:
    """Fail-closed live broker. It never requests or returns credential values."""

    def project(self, engagement: Engagement, registration: dict[str, Any]
                ) -> tuple[IdentityHandle, ...]:
        raise AuthorizationError(
            "live custody broker is not connected; no credential was requested"
        )

    def revoke(self, lease_id: str) -> None:
        raise AuthorizationError(
            f"live custody broker is not connected; lease {lease_id} was not created"
        )


class LocalBroker:
    """In-process projection for fixture-e2. Secret bytes never leave this object."""

    def __init__(self) -> None:
        self._secrets: dict[str, bytes] = {}
        self._leases: dict[str, tuple[str, str]] = {}

    def project(self, engagement: Engagement, registration: dict[str, Any]
                ) -> tuple[IdentityHandle, ...]:
        if registration.get("applicability") != "applicable":
            raise AuthorizationError("broker will not project identities for a non-applicable target")
        count = int(registration["identities"]["count"])
        lease_id = os.urandom(8).hex()
        ttl = min(int(engagement.raw["credential_max_ttl_seconds"]), STANDING_E2_TTL)
        expires_at = (datetime.now(UTC) + timedelta(seconds=ttl)).isoformat().replace("+00:00", "Z")
        handles: list[IdentityHandle] = []
        if count >= 1:
            role = "owner" if count == 2 else "runtime"
            path = f"/var/run/secrets/whitehat/{lease_id}/{'token-b' if count == 2 else 'runtime'}"
            handles.append(IdentityHandle(role, path, lease_id, expires_at))
            self._secrets[path] = os.urandom(32)
        if count == 2:
            path = f"/var/run/secrets/whitehat/{lease_id}/token-a"
            handles.append(IdentityHandle("attacker", path, lease_id, expires_at))
            self._secrets[path] = os.urandom(32)
        self._leases[lease_id] = tuple(handle.mount_path for handle in handles)
        return tuple(handles)

    def revoke(self, lease_id: str) -> None:
        paths = self._leases.pop(lease_id, None)
        if paths is None:
            raise AuthorizationError("lease is unknown; cleanup cannot be proven")
        for path in paths:
            self._secrets.pop(path, None)


@dataclass
class PlaneLease:
    engagement: Engagement
    registration: dict[str, Any]
    identities: tuple[IdentityHandle, ...]
    watcher: RateWatcher
    lease_id: str
    broker_name: str

    def permits(self, *, technique: str, route: str) -> None:
        self.engagement.permits(technique=technique, route=route)
        if route not in self.registration["routes"]:
            raise AuthorizationError(f"route not registered: {route}")


def default_broker(engagement: Engagement) -> CredentialBroker:
    if (engagement.raw.get("environment") == "fixture"
            and engagement.raw.get("approval_class") in FIXTURE_CLASSES):
        return LocalBroker()
    return UnconnectedCustodyBroker()


def admit(*, engagement: Engagement, registration: dict[str, Any],
          broker: CredentialBroker | None = None,
          kill_switch: KillSwitch | None = None,
          now: datetime | None = None,
          retired: set[str] | None = None) -> PlaneLease:
    engagement.validate(now=now)
    switch = kill_switch or KillSwitch()
    if switch.engaged():
        raise AuthorizationError("kill switch is engaged")
    engagement_id = engagement.raw["engagement_id"]
    if engagement_id in (retired if retired is not None else retired_ids()):
        raise AuthorizationError("engagement id is retired and must not be reused")
    if registration.get("applicability") != "applicable":
        raise AuthorizationError(
            f"target is {registration.get('applicability', 'unregistered')}"
        )
    target_id = engagement.raw.get("target_id")
    if not target_id:
        raise AuthorizationError("engagement missing target_id for plane admission")
    if target_id != registration["target_id"]:
        raise AuthorizationError("engagement target_id does not match registration")
    approval_class = engagement.raw.get("approval_class")
    allowed = APPROVAL_CLASSES.get(approval_class)
    if allowed is None:
        raise AuthorizationError("engagement missing or unknown approval_class")
    if approval_class not in registration["approval_classes"]:
        raise AuthorizationError("approval_class is not registered for this target")
    for technique in engagement.raw["techniques"]:
        if technique not in allowed:
            raise AuthorizationError(
                f"technique {technique} not in approval class {approval_class}"
            )
    for route in engagement.raw["routes"]:
        if route not in registration["routes"]:
            raise AuthorizationError(f"route not registered: {route}")
    _enforce_class_envelope(approval_class, engagement)
    selected = broker or default_broker(engagement)
    count = int(registration["identities"]["count"])
    if count:
        identities = selected.project(engagement, registration)
        if len(identities) != count:
            raise AuthorizationError(f"broker must project exactly {count} identities")
        expected = {2: {"owner", "attacker"}, 1: {"runtime"}}[count]
        if {handle.role for handle in identities} != expected:
            raise AuthorizationError(f"broker must project {sorted(expected)} identities")
        if any(hasattr(handle, "secret") and getattr(handle, "secret") for handle in identities):
            raise AuthorizationError("broker exposed a credential value")
        lease_id = identities[0].lease_id
    else:
        if approval_class not in FIXTURE_CLASSES:
            selected.project(engagement, registration)
        identities = ()
        lease_id = os.urandom(8).hex()
        if hasattr(selected, "_leases"):
            selected._leases[lease_id] = ()
    watcher = RateWatcher(
        per_minute=engagement.raw["rate_limit_per_minute"],
        max_concurrency=engagement.raw["max_concurrency"],
        max_requests=int(engagement.raw.get("maximum_requests")
                         or engagement.raw["rate_limit_per_minute"]),
    )
    return PlaneLease(
        engagement=engagement,
        registration=registration,
        identities=identities,
        watcher=watcher,
        lease_id=lease_id,
        broker_name=type(selected).__name__,
    )


def _enforce_class_envelope(approval_class: str, engagement: Engagement) -> None:
    environment = engagement.raw["environment"]
    if approval_class in FIXTURE_CLASSES:
        if environment != "fixture":
            raise AuthorizationError(f"{approval_class} requires environment=fixture")
        return
    if environment == "fixture":
        raise AuthorizationError(f"{approval_class} cannot use the fixture environment")
    if engagement.raw.get("plane_namespace") != PLANE_NAMESPACE:
        raise AuthorizationError(f"{approval_class} requires the whitehat plane namespace")
    if not engagement.raw.get("runner_image_digest"):
        raise AuthorizationError(f"{approval_class} requires a pinned runner image digest")
    if engagement.raw["credential_max_ttl_seconds"] > STANDING_E2_TTL:
        raise AuthorizationError("credential TTL must be at most 900 seconds")
    if approval_class == "e3" and not engagement.raw.get("database"):
        raise AuthorizationError("live E3 requires a named database")
    if approval_class == "capacity" and engagement.raw.get("aggressor_ceiling") is None:
        raise AuthorizationError("live capacity requires an aggressor_ceiling")


def cleanup(lease: PlaneLease, broker: CredentialBroker) -> dict[str, str]:
    broker.revoke(lease.lease_id)
    return {
        "fixtures": "delete only recorded fixture_ids",
        "credential_revocation": "revoked",
        "lease_id": lease.lease_id,
    }