"""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()