bot-bottle/bot_bottle/orchestrator/broker.py

"""Launch broker contract (PRD 0070).

A VM/container can't spawn its own host-network siblings, so the
orchestrator brokers agent launches through a small privileged component.
The request it sends is:

  * **structured** — static flags + ids only (bottle id, pool slot, a
    content-addressed image ref), never a free-form path or argv, so a
    request can't be coerced into launching an arbitrary payload; and
  * **signed** — wrapped as a compact JWS/JWT the broker verifies before
    acting, so a *compromised co-located component* (an agent-facing
    gateway, say) can't forge a launch. This is the concrete form of the
    "structured requests only" rule in the PRD security review.

Signing is **HS256** over a secret shared by the orchestrator (signer) and
the broker (verifier) — appropriate here because both are trusted host
components provisioned together; the secret is exactly what an
agent-facing component does not have. (Stdlib only — the project takes no
runtime deps; a cross-host deployment could swap in an asymmetric alg.)
"""

from __future__ import annotations

import abc
import base64
import hashlib
import hmac
import json
import secrets
import time
from dataclasses import dataclass
from typing import Protocol

_JWT_HEADER = {"alg": "HS256", "typ": "JWT"}

# The closed op vocabulary, split by kind (PRD "ids + static flags" rule):
#   * mutation ops act on ONE bottle — they carry its id (+ static launch flags);
#   * query ops enumerate host state — they carry NO ids or flags at all.
# Each new host-privileged op is added here deliberately; `verify_request`
# enforces the per-kind shape so the privileged surface can't drift.
_MUTATION_OPS = ("launch", "teardown")
_QUERY_OPS = ("list_live",)
_ALLOWED_OPS = _MUTATION_OPS + _QUERY_OPS

# Every claim key a signed request may carry — the request fields plus the two
# per-signature envelope claims. A strict allow-list (open question 1, resolved
# "yes"): an unknown key is rejected outright, so the schema can't silently widen
# as ops are added.
_ALLOWED_CLAIMS = frozenset(
    {"op", "bottle_id", "source_ip", "image_ref", "slot", "jti", "iat"}
)


class BrokerAuthError(Exception):
    """A broker request failed provenance or schema verification —
    bad/absent signature, malformed token, or a payload that doesn't match
    the fixed launch-request shape. Fail-closed: the broker must not act.

    A **definite** negative: nothing was launched, so a caller may safely roll
    back as if the op never happened."""


class BrokerUnavailableError(Exception):
    """A brokered request could not be carried to a verdict: the broker (or the
    wire to it) was unreachable, timed out, or dropped the response.

    Crucially **ambiguous** — unlike `BrokerAuthError`, the op MAY already have
    taken effect on the backend before the response was lost, so a caller must
    NOT assume it did nothing (e.g. must not roll a registry row back as if no
    launch happened, which would orphan a running container). Only the in-process
    brokers never raise this; the out-of-process `BrokerClient` does."""


@dataclass(frozen=True)
class BrokerRequest:
    """A structured, un-coercible broker request. Ids + static flags only —
    `image_ref` is a content-addressed id, never a path/argv.

    A **mutation** op (`launch`/`teardown`) names the bottle it acts on in
    `bottle_id` (plus launch's static flags). A **query** op (`list_live`) acts
    on no single bottle, so it carries nothing but its `op` name — `bottle_id`
    and every flag stay empty, and `verify_request` rejects a query that smuggles
    any in. `LaunchRequest` is kept as an alias for the launch/teardown callers."""

    op: str  # one of _ALLOWED_OPS
    bottle_id: str = ""
    source_ip: str = ""
    image_ref: str = ""
    slot: int | None = None


# Historical name — the request type predates the query ops. Kept so the launch
# path (`OrchestratorCore`, `DockerBroker`) reads unchanged.
LaunchRequest = BrokerRequest


# --- minimal JWS/JWT (HS256), stdlib only ----------------------------------

def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(text: str) -> bytes:
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


def _mac(signing_input: str, secret: bytes) -> str:
    digest = hmac.new(secret, signing_input.encode("ascii"), hashlib.sha256).digest()
    return _b64url(digest)


def sign_request(req: LaunchRequest, secret: bytes) -> str:
    """Serialize `req` to signed-JWT form. Adds a per-signature `jti`
    (replay id) and `iat` (issued-at)."""
    claims: dict[str, object] = {
        "op": req.op,
        "bottle_id": req.bottle_id,
        "source_ip": req.source_ip,
        "image_ref": req.image_ref,
        "slot": req.slot,
        "jti": secrets.token_hex(8),
        "iat": int(time.time()),
    }
    header = _b64url(json.dumps(_JWT_HEADER, sort_keys=True, separators=(",", ":")).encode())
    payload = _b64url(json.dumps(claims, sort_keys=True, separators=(",", ":")).encode())
    signing_input = f"{header}.{payload}"
    return f"{signing_input}.{_mac(signing_input, secret)}"


def verify_request(token: str, secret: bytes) -> LaunchRequest:
    """Verify the signature and the request shape; return the parsed
    request. Raises `BrokerAuthError` on any failure (fail-closed)."""
    parts = token.split(".")
    if len(parts) != 3:
        raise BrokerAuthError("malformed token")
    header_b, payload_b, sig = parts
    if not hmac.compare_digest(_mac(f"{header_b}.{payload_b}", secret), sig):
        raise BrokerAuthError("bad signature")
    try:
        header = json.loads(_b64url_decode(header_b))
        claims = json.loads(_b64url_decode(payload_b))
    except (ValueError, json.JSONDecodeError) as e:
        raise BrokerAuthError("malformed token payload") from e
    if not isinstance(header, dict) or header.get("alg") != "HS256":
        raise BrokerAuthError("unexpected header/alg")
    if not isinstance(claims, dict):
        raise BrokerAuthError("claims are not an object")
    # Strict schema (open question 1): a signed request may carry ONLY the known
    # claim keys, so the privileged surface can't widen by smuggling an extra
    # claim past the fixed field checks below.
    unknown = set(claims) - _ALLOWED_CLAIMS
    if unknown:
        raise BrokerAuthError(f"unknown claim(s): {', '.join(sorted(unknown))}")

    op = claims.get("op")
    if not isinstance(op, str) or op not in _ALLOWED_OPS:
        raise BrokerAuthError("request does not match the broker-request schema")
    bottle_id = claims.get("bottle_id", "")
    source_ip = claims.get("source_ip", "")
    image_ref = claims.get("image_ref", "")
    slot = claims.get("slot")
    if (
        not isinstance(bottle_id, str)
        or not isinstance(source_ip, str) or not isinstance(image_ref, str)
        or not (slot is None or isinstance(slot, int))
    ):
        raise BrokerAuthError("request does not match the broker-request schema")
    # Per-kind shape: a mutation names exactly one bottle; a query names none and
    # carries no flags (the "no arguments" rule for `list_live`). Enforcing both
    # halves keeps `bottle_id`/flags from being a coercible field on a query op.
    if op in _MUTATION_OPS:
        if not bottle_id:
            raise BrokerAuthError(f"{op} requires a bottle_id")
    elif bottle_id or source_ip or image_ref or slot is not None:
        raise BrokerAuthError(f"{op} is a query op and takes no ids or flags")
    return BrokerRequest(
        op=op, bottle_id=bottle_id, source_ip=source_ip, image_ref=image_ref, slot=slot
    )


# --- the broker itself ------------------------------------------------------

class SubmitBroker(Protocol):
    """The broker surface `OrchestratorCore` depends on. `submit` verifies a
    signed *mutation* token and performs its op, returning the verified request;
    `list_live` verifies a signed `list_live` token and returns the source IPs of
    the bottles the backend currently has running (reconcile's live set). Both the
    in-process `LaunchBroker` and the out-of-process `BrokerClient` (which relays
    the token to the host control server) satisfy it structurally, so the core is
    unchanged whether the backend is local or a real host service."""

    def submit(self, token: str) -> LaunchRequest: ...

    def list_live(self, token: str) -> list[str]: ...


class LaunchBroker(abc.ABC):
    """Verifies a signed request came from the orchestrator, then performs
    the backend-native launch/teardown. Subclasses implement `_launch` /
    `_teardown`; verification is shared and fail-closed."""

    def __init__(self, secret: bytes) -> None:
        self._secret = secret

    def submit(self, token: str) -> LaunchRequest:
        """Verify `token` and perform its mutation op. Returns the verified
        request; raises `BrokerAuthError` if provenance/shape fails, or if the
        token is a non-mutation op (a `list_live` token routed here by mistake)."""
        req = verify_request(token, self._secret)
        if req.op == "launch":
            self._launch(req)
        elif req.op == "teardown":
            self._teardown(req)
        else:
            raise BrokerAuthError(f"{req.op} is not a submit op")
        return req

    def list_live(self, token: str) -> list[str]:
        """Verify a `list_live` `token` and return the source IPs of the bottles
        the backend currently has running. Raises `BrokerAuthError` on bad
        provenance/shape or a non-`list_live` op; a backend enumeration failure
        is converted to `BrokerUnavailableError` — the single "live set could not
        be determined" signal reconcile catches to skip the sweep rather than
        reaping healthy rows against a partial set (the out-of-process
        `BrokerClient` raises the same on an enumeration 502 / no response)."""
        req = verify_request(token, self._secret)
        if req.op != "list_live":
            raise BrokerAuthError(f"{req.op} is not a list_live op")
        try:
            return list(self._list_live())
        except Exception as e:  # noqa: BLE001 — any enumeration failure means the
            # live set is unknown; reconcile must skip, never reap against it.
            raise BrokerUnavailableError(f"list_live enumeration failed: {e}") from e

    @abc.abstractmethod
    def _launch(self, req: LaunchRequest) -> None:
        ...

    @abc.abstractmethod
    def _teardown(self, req: LaunchRequest) -> None:
        ...

    @abc.abstractmethod
    def _list_live(self) -> list[str]:
        """Every running bottle's source IP, backend-native. Must raise (not
        return a partial list) if the live set can't be determined
        authoritatively, so reconcile skips rather than reaping healthy rows."""
        ...


class StubBroker(LaunchBroker):
    """Dev-harness broker: records verified requests without launching
    anything. Exercises the full sign -> verify -> act contract in-process.

    Its live set (`_list_live`) is, by default, every bottle it recorded a
    launch for and no teardown since — the in-process analogue of enumerating
    the backend. Tests that need to simulate a bottle dying out from under the
    registry (the case reconcile exists for) set `live_source_ips` explicitly to
    override that derived set."""

    def __init__(self, secret: bytes) -> None:
        super().__init__(secret)
        self.launched: list[LaunchRequest] = []
        self.torn_down: list[LaunchRequest] = []
        # When not None, the exact live set `_list_live` reports — lets a test
        # say "only these IPs are still up" regardless of what was launched.
        self.live_source_ips: list[str] | None = None

    def _launch(self, req: LaunchRequest) -> None:
        self.launched.append(req)

    def _teardown(self, req: LaunchRequest) -> None:
        self.torn_down.append(req)

    def _list_live(self) -> list[str]:
        if self.live_source_ips is not None:
            return list(self.live_source_ips)
        torn = {r.bottle_id for r in self.torn_down}
        return [r.source_ip for r in self.launched
                if r.bottle_id not in torn and r.source_ip]


__all__ = [
    "BrokerAuthError",
    "BrokerUnavailableError",
    "BrokerRequest",
    "LaunchRequest",
    "SubmitBroker",
    "LaunchBroker",
    "StubBroker",
    "sign_request",
    "verify_request",
]