bot-bottle/tests/unit/test_orchestrator_broker.py

"""Unit tests for the orchestrator launch broker (PRD 0070)."""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import unittest

from bot_bottle.orchestrator.broker import (
    BrokerAuthError,
    LaunchRequest,
    StubBroker,
    sign_request,
    verify_request,
)


def _sign_claims(claims: dict[str, object], secret: bytes) -> str:
    """Sign an arbitrary claims dict (bypassing `sign_request`'s fixed fields)
    so a test can craft off-schema payloads — e.g. an extra claim key, or a
    query op carrying ids — with a *valid* signature, isolating the schema check
    from the provenance check."""
    def b64(data: bytes) -> str:
        return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
    header = b64(json.dumps({"alg": "HS256", "typ": "JWT"}, sort_keys=True,
                            separators=(",", ":")).encode())
    payload = b64(json.dumps(claims, sort_keys=True, separators=(",", ":")).encode())
    signing_input = f"{header}.{payload}"
    sig = b64(hmac.new(secret, signing_input.encode("ascii"), hashlib.sha256).digest())
    return f"{signing_input}.{sig}"


class TestSignVerify(unittest.TestCase):
    def setUp(self) -> None:
        self.secret = secrets.token_bytes(16)

    def test_launch_round_trip(self) -> None:
        req = LaunchRequest(
            op="launch", bottle_id="b1", source_ip="10.243.0.1",
            image_ref="sha256:abc", slot=3,
        )
        self.assertEqual(req, verify_request(sign_request(req, self.secret), self.secret))

    def test_teardown_round_trip(self) -> None:
        req = LaunchRequest(op="teardown", bottle_id="b1")
        got = verify_request(sign_request(req, self.secret), self.secret)
        self.assertEqual(("teardown", "b1"), (got.op, got.bottle_id))

    def test_wrong_secret_rejected(self) -> None:
        token = sign_request(LaunchRequest(op="launch", bottle_id="b1"), self.secret)
        with self.assertRaises(BrokerAuthError):
            verify_request(token, secrets.token_bytes(16))

    def test_tampered_payload_rejected(self) -> None:
        token = sign_request(LaunchRequest(op="launch", bottle_id="b1"), self.secret)
        header, payload, sig = token.split(".")
        flipped = payload[:-1] + ("A" if payload[-1] != "A" else "B")
        with self.assertRaises(BrokerAuthError):
            verify_request(f"{header}.{flipped}.{sig}", self.secret)

    def test_malformed_tokens_rejected(self) -> None:
        for bad in ("", "a.b", "a.b.c.d", "not-a-token"):
            with self.assertRaises(BrokerAuthError):
                verify_request(bad, self.secret)

    def test_unknown_op_rejected_despite_valid_signature(self) -> None:
        # A correctly-signed token whose op isn't in the allow-list must
        # still be refused — the schema guard is independent of provenance.
        token = sign_request(LaunchRequest(op="rm-rf", bottle_id="b1"), self.secret)
        with self.assertRaises(BrokerAuthError):
            verify_request(token, self.secret)

    def test_list_live_round_trip(self) -> None:
        req = LaunchRequest(op="list_live")
        got = verify_request(sign_request(req, self.secret), self.secret)
        self.assertEqual("list_live", got.op)
        self.assertEqual("", got.bottle_id)

    def test_mutation_without_bottle_id_rejected(self) -> None:
        # A launch/teardown must name its bottle even with a valid signature.
        for op in ("launch", "teardown"):
            token = _sign_claims(
                {"op": op, "bottle_id": "", "jti": "x", "iat": 1}, self.secret)
            with self.assertRaises(BrokerAuthError):
                verify_request(token, self.secret)

    def test_query_op_carrying_ids_or_flags_rejected(self) -> None:
        # list_live is "no arguments" — a signed one that smuggles a bottle_id /
        # source_ip / image_ref / slot is off-schema and refused.
        for extra in (
            {"bottle_id": "b1"}, {"source_ip": "10.0.0.1"},
            {"image_ref": "img"}, {"slot": 2},
        ):
            claims: dict[str, object] = {"op": "list_live", "jti": "x", "iat": 1, **extra}
            with self.assertRaises(BrokerAuthError):
                verify_request(_sign_claims(claims, self.secret), self.secret)

    def test_unknown_claim_key_rejected(self) -> None:
        # Strict schema (open question 1): a well-signed token with an extra claim
        # key is refused, so the surface can't widen past the fixed fields.
        token = _sign_claims(
            {"op": "launch", "bottle_id": "b1", "cmd": "rm -rf /", "jti": "x", "iat": 1},
            self.secret,
        )
        with self.assertRaises(BrokerAuthError):
            verify_request(token, self.secret)


class TestStubBroker(unittest.TestCase):
    def setUp(self) -> None:
        self.secret = secrets.token_bytes(16)
        self.broker = StubBroker(self.secret)

    def test_submit_launch_records_verified_request(self) -> None:
        req = LaunchRequest(op="launch", bottle_id="b1", source_ip="10.243.0.1")
        got = self.broker.submit(sign_request(req, self.secret))
        self.assertEqual(req, got)
        self.assertEqual(["b1"], [r.bottle_id for r in self.broker.launched])
        self.assertEqual([], self.broker.torn_down)

    def test_submit_teardown_records(self) -> None:
        self.broker.submit(
            sign_request(LaunchRequest(op="teardown", bottle_id="b1"), self.secret)
        )
        self.assertEqual(["b1"], [r.bottle_id for r in self.broker.torn_down])

    def test_submit_forged_token_is_fail_closed(self) -> None:
        forged = sign_request(
            LaunchRequest(op="launch", bottle_id="b1"), secrets.token_bytes(16)
        )
        with self.assertRaises(BrokerAuthError):
            self.broker.submit(forged)
        self.assertEqual([], self.broker.launched)  # nothing acted on

    def test_list_live_derives_from_launches(self) -> None:
        # Default live set: launched, minus torn-down, by source IP.
        self.broker.submit(
            sign_request(LaunchRequest(op="launch", bottle_id="b1", source_ip="10.0.0.1"),
                         self.secret))
        self.broker.submit(
            sign_request(LaunchRequest(op="launch", bottle_id="b2", source_ip="10.0.0.2"),
                         self.secret))
        self.broker.submit(
            sign_request(LaunchRequest(op="teardown", bottle_id="b1", source_ip="10.0.0.1"),
                         self.secret))
        live = self.broker.list_live(
            sign_request(LaunchRequest(op="list_live"), self.secret))
        self.assertEqual(["10.0.0.2"], live)

    def test_list_live_override(self) -> None:
        self.broker.live_source_ips = ["10.0.0.9"]
        live = self.broker.list_live(
            sign_request(LaunchRequest(op="list_live"), self.secret))
        self.assertEqual(["10.0.0.9"], live)

    def test_submit_rejects_a_list_live_token(self) -> None:
        # The verbs don't cross: a query token routed to submit is fail-closed.
        with self.assertRaises(BrokerAuthError):
            self.broker.submit(sign_request(LaunchRequest(op="list_live"), self.secret))

    def test_list_live_rejects_a_mutation_token(self) -> None:
        with self.assertRaises(BrokerAuthError):
            self.broker.list_live(
                sign_request(LaunchRequest(op="launch", bottle_id="b1"), self.secret))


if __name__ == "__main__":
    unittest.main()