From 1702664b8192945be3c47e96c11cc07dd262f7d6 Mon Sep 17 00:00:00 2001 From: didericis Date: Mon, 13 Jul 2026 13:00:04 -0400 Subject: [PATCH] =?UTF-8?q?feat(orchestrator):=20slice=201=20=E2=80=94=20r?= =?UTF-8?q?egistry=20+=20attribution=20+=20HTTP=20control=20plane=20(#352)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First implementation slice of PRD 0070, the backend-neutral consolidation core as a plain-process dev-harness (no VM packaging yet): * orchestrator/registry.py — SQLite (WAL) runtime-state store on the existing DbStore/TableMigrations base. Live bottle registry keyed by source IP + per-bottle identity token, with fail-closed attribution: a request resolves to a bottle only when its source IP AND identity token both match exactly one active record (unknown/ambiguous IP, empty token, or token mismatch all deny). Tokens are 256-bit urandom. * orchestrator/control_plane.py — the HTTP control plane (the universal transport chosen in 0070): register / deregister / list / attribute / health. Routing is a pure dispatch() so it is socket-free testable; Handler/ControlPlaneServer/make_server are a thin stdlib adapter. register/deregister are the live-reload path; listing redacts tokens. * orchestrator/__main__.py — `python -m bot_bottle.orchestrator` harness. Launch/teardown, the launch broker, and the egress/git/supervise data plane come in later slices. 24 unit tests (attribution matrix, persistence, dispatch, one real-socket round-trip). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01WBMWTEtQdJ4W5UrWuLHCck --- bot_bottle/orchestrator/__init__.py | 30 +++ bot_bottle/orchestrator/__main__.py | 52 +++++ bot_bottle/orchestrator/control_plane.py | 154 +++++++++++++ bot_bottle/orchestrator/registry.py | 212 ++++++++++++++++++ tests/unit/test_orchestrator_control_plane.py | 150 +++++++++++++ tests/unit/test_orchestrator_registry.py | 102 +++++++++ 6 files changed, 700 insertions(+) create mode 100644 bot_bottle/orchestrator/__init__.py create mode 100644 bot_bottle/orchestrator/__main__.py create mode 100644 bot_bottle/orchestrator/control_plane.py create mode 100644 bot_bottle/orchestrator/registry.py create mode 100644 tests/unit/test_orchestrator_control_plane.py create mode 100644 tests/unit/test_orchestrator_registry.py diff --git a/bot_bottle/orchestrator/__init__.py b/bot_bottle/orchestrator/__init__.py new file mode 100644 index 0000000..b04bab4 --- /dev/null +++ b/bot_bottle/orchestrator/__init__.py @@ -0,0 +1,30 @@ +"""Per-host orchestrator service (PRD 0070). + +A single persistent per-host service that will run the sidecar functions +(egress / git-gate / supervise), coordinate with the console, and broker +agent launches. This package is being built bottom-up, starting with the +backend-neutral "consolidation core" that needs no VM packaging: + + * `registry` — the SQLite runtime-state store + fail-closed + attribution (source IP + per-bottle identity token). + * `control_plane` — the HTTP control-plane RPC (register / deregister / + list / attribute / health), with live reload. + +Launch/teardown, the launch broker, and the actual egress/git/supervise +data plane land in later slices once this core is proven (see the PRD's +sequencing: plain-process dev-harness -> docker orchestrator -> firecracker). +""" + +from __future__ import annotations + +from .registry import BottleRecord, RegistryStore, new_identity_token +from .control_plane import ControlPlaneServer, dispatch, make_server + +__all__ = [ + "BottleRecord", + "RegistryStore", + "new_identity_token", + "ControlPlaneServer", + "dispatch", + "make_server", +] diff --git a/bot_bottle/orchestrator/__main__.py b/bot_bottle/orchestrator/__main__.py new file mode 100644 index 0000000..7a2fd9a --- /dev/null +++ b/bot_bottle/orchestrator/__main__.py @@ -0,0 +1,52 @@ +"""Run the orchestrator control plane as a plain process (PRD 0070 dev-harness). + + python -m bot_bottle.orchestrator [--host H] [--port P] [--db PATH] + +The PRD sequences the orchestrator as a plain-process dev-harness first, so +the consolidation core (registry + attribution + HTTP control plane + live +reload) can be exercised with fast iteration, decoupled from any VM / +container packaging. Wrapping this exact service in a backend-native unit +(docker container, then Firecracker VM) comes later. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from .. import log +from .control_plane import make_server +from .registry import RegistryStore, default_db_path + + +def main(argv: list[str] | None = None) -> int: + """Parse args, migrate the registry, and serve the control plane.""" + parser = argparse.ArgumentParser(prog="bot_bottle.orchestrator") + parser.add_argument("--host", default="127.0.0.1", help="bind address") + parser.add_argument("--port", type=int, default=8080, help="bind port (0 = ephemeral)") + parser.add_argument( + "--db", type=Path, default=None, + help=f"registry DB path (default: {default_db_path()})", + ) + args = parser.parse_args(argv) + + registry = RegistryStore(args.db) + registry.migrate() + + server = make_server(registry, host=args.host, port=args.port) + bound_host, bound_port = server.server_address[0], server.server_address[1] + log.info( + "orchestrator control plane listening", + context={"host": bound_host, "port": bound_port, "db": str(registry.db_path)}, + ) + try: + server.serve_forever() + except KeyboardInterrupt: + log.info("orchestrator shutting down") + finally: + server.server_close() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bot_bottle/orchestrator/control_plane.py b/bot_bottle/orchestrator/control_plane.py new file mode 100644 index 0000000..772eae2 --- /dev/null +++ b/bot_bottle/orchestrator/control_plane.py @@ -0,0 +1,154 @@ +"""Orchestrator HTTP control plane (PRD 0070). + +The backend-agnostic control-plane RPC (CLI / console -> orchestrator) over +**HTTP** — the universal transport chosen in 0070 (works on every host; no +vsock / unix-socket portability caveats). Endpoints mutate the live +registry, so register/deregister are the *live-reload* path — no restart: + + GET /health -> 200 {"status": "ok"} + GET /bottles -> 200 {"bottles": [ , ... ]} + POST /bottles -> 201 {"bottle_id", "identity_token"} + body: {"source_ip", ["bottle_id"], ["metadata"]} + DELETE /bottles/ -> 200 {"deregistered": true} | 404 + POST /attribute -> 200 {"bottle_id"} | 403 + body: {"source_ip", "identity_token"} + +Routing/handling is the pure function `dispatch()` so it is unit-testable +without a socket; `Handler` / `ControlPlaneServer` / `make_server` are a +thin stdlib adapter around it. + +Note the listing redacts identity tokens — they are never returned except +once, to the caller that registers the bottle. +""" + +from __future__ import annotations + +import http.server +import json +import os +import socketserver +import typing +from urllib.parse import urlsplit + +from .registry import RegistryStore + +# JSON body payload type (parsed request / rendered response). +Json = dict[str, object] + + +def _parse_json_object(body: bytes) -> Json: + """Parse a JSON object body. Raises ValueError for non-objects / bad JSON.""" + if not body: + return {} + obj = json.loads(body) # raises json.JSONDecodeError (a ValueError) + if not isinstance(obj, dict): + raise ValueError("request body must be a JSON object") + return obj + + +def dispatch( # pylint: disable=too-many-return-statements + registry: RegistryStore, method: str, path: str, body: bytes +) -> tuple[int, Json]: + """Route one control-plane request to a (status, payload) pair. Pure — + no I/O beyond the registry — so it is fully testable without a socket.""" + route = urlsplit(path).path.rstrip("/") or "/" + + if method == "GET" and route == "/health": + return 200, {"status": "ok"} + + if method == "GET" and route == "/bottles": + return 200, {"bottles": [r.redacted() for r in registry.all()]} + + if method == "POST" and route == "/bottles": + try: + data = _parse_json_object(body) + except ValueError as e: + return 400, {"error": f"invalid JSON: {e}"} + source_ip = data.get("source_ip") + if not isinstance(source_ip, str) or not source_ip: + return 400, {"error": "source_ip (string) is required"} + bottle_id = data.get("bottle_id") + metadata = data.get("metadata") + rec = registry.register( + source_ip, + bottle_id=bottle_id if isinstance(bottle_id, str) else None, + metadata=metadata if isinstance(metadata, str) else "", + ) + return 201, {"bottle_id": rec.bottle_id, "identity_token": rec.identity_token} + + if method == "DELETE" and route.startswith("/bottles/"): + bottle_id = route[len("/bottles/"):] + if registry.deregister(bottle_id): + return 200, {"deregistered": True} + return 404, {"error": "no such bottle"} + + if method == "POST" and route == "/attribute": + try: + data = _parse_json_object(body) + except ValueError as e: + return 400, {"error": f"invalid JSON: {e}"} + source_ip = data.get("source_ip") + token = data.get("identity_token") + if not isinstance(source_ip, str) or not isinstance(token, str): + return 400, {"error": "source_ip and identity_token (strings) required"} + rec = registry.attribute(source_ip, token) + if rec is None: + return 403, {"error": "unattributed"} + return 200, {"bottle_id": rec.bottle_id} + + return 404, {"error": "not found"} + + +class Handler(http.server.BaseHTTPRequestHandler): + """Thin stdlib adapter: read the body, call `dispatch`, write JSON.""" + + # Quiet by default (the orchestrator has its own logging); opt back into + # stdlib access logging with BOT_BOTTLE_ORCHESTRATOR_DEBUG. + def log_message(self, format: str, *args: typing.Any) -> None: # noqa: A002 + if os.environ.get("BOT_BOTTLE_ORCHESTRATOR_DEBUG"): + super().log_message(format, *args) + + def _serve(self, method: str) -> None: + """Read the request body, dispatch it, and write the JSON reply.""" + server = self.server + assert isinstance(server, ControlPlaneServer) + length = int(self.headers.get("Content-Length") or 0) + body = self.rfile.read(length) if length > 0 else b"" + status, payload = dispatch(server.registry, method, self.path, body) + data = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def do_GET(self) -> None: + self._serve("GET") + + def do_POST(self) -> None: + self._serve("POST") + + def do_DELETE(self) -> None: + self._serve("DELETE") + + +class ControlPlaneServer(socketserver.ThreadingMixIn, http.server.HTTPServer): + """Threading HTTP server that carries the registry for its handlers.""" + + daemon_threads = True + allow_reuse_address = True + + def __init__(self, address: tuple[str, int], registry: RegistryStore) -> None: + self.registry = registry + super().__init__(address, Handler) + + +def make_server( + registry: RegistryStore, host: str = "127.0.0.1", port: int = 0 +) -> ControlPlaneServer: + """Build (but do not start) a control-plane server. `port=0` binds an + ephemeral port — read `server.server_address` for the actual one.""" + return ControlPlaneServer((host, port), registry) + + +__all__ = ["dispatch", "Handler", "ControlPlaneServer", "make_server", "Json"] diff --git a/bot_bottle/orchestrator/registry.py b/bot_bottle/orchestrator/registry.py new file mode 100644 index 0000000..df35a12 --- /dev/null +++ b/bot_bottle/orchestrator/registry.py @@ -0,0 +1,212 @@ +"""Orchestrator bottle registry — the per-host runtime-state store (PRD 0070). + +SQLite-backed (WAL) registry mapping each live bottle to its source IP and +per-bottle identity token, plus the **fail-closed attribution** the data +plane relies on: a request is attributed to a bottle only when its source +IP *and* its identity token both match a single active record. + +This is the "runtime state" tier of PRD 0070 (leases / approvals / +registry) — deliberately NOT config (which stays declarative under +`~/.bot-bottle/`) and NOT the build-time constants (a flat file). It is the +source of truth the orchestrator sweeps on restart to re-adopt running +bottles, so it lives in a durable store rather than process memory. + +Attribution combines two independent signals, and needs both: + + * source IP — the network-layer invariant (on Firecracker the `/31` TAP + + nft make it unspoofable by construction; weaker on other backends). + * identity token — an application-layer per-bottle secret, defence in + depth that doesn't lean on network anti-spoof. A bottle can only ever + prove it is *itself* (it can't learn another bottle's token), so a + hostile agent gains nothing by presenting it. +""" + +from __future__ import annotations + +import hmac +import secrets +import sqlite3 +import time +from dataclasses import dataclass +from pathlib import Path + +from ..db_store import DbStore +from ..migrations import TableMigrations + +# 256 bits of urandom, URL-safe — unguessable per-bottle identity token. +IDENTITY_TOKEN_BYTES = 32 + + +def new_identity_token() -> str: + """A fresh per-bottle identity token (PRD 0070 attribution defence).""" + return secrets.token_urlsafe(IDENTITY_TOKEN_BYTES) + + +def default_db_path() -> Path: + """Host location for the orchestrator's runtime-state DB. Runtime state + (not config), so it lives under the cache dir like the pool state.""" + return Path.home() / ".cache" / "bot-bottle" / "orchestrator" / "registry.db" + + +@dataclass(frozen=True) +class BottleRecord: + """One live bottle in the registry.""" + + bottle_id: str + source_ip: str + identity_token: str + state: str = "active" + created_at: float = 0.0 + # Opaque JSON blob for forward-compat (policy refs, pool slot, ...) — + # the registry doesn't interpret it, so new fields need no migration. + metadata: str = "" + + def redacted(self) -> dict[str, object]: + """Public view for listing — never exposes the identity token.""" + return { + "bottle_id": self.bottle_id, + "source_ip": self.source_ip, + "state": self.state, + "created_at": self.created_at, + "metadata": self.metadata, + } + + +_MIGRATIONS = TableMigrations( + "orchestrator_registry", + [ + # v1 — the live bottle registry. + """ + CREATE TABLE IF NOT EXISTS orchestrator_bottles ( + bottle_id TEXT PRIMARY KEY, + source_ip TEXT NOT NULL, + identity_token TEXT NOT NULL, + state TEXT NOT NULL DEFAULT 'active', + created_at REAL NOT NULL, + metadata TEXT NOT NULL DEFAULT '' + ) + """, + # source_ip is the data-plane attribution key — index it, and make + # the "one active bottle per source IP" check cheap. + "CREATE INDEX IF NOT EXISTS idx_orchestrator_bottles_source_ip " + "ON orchestrator_bottles (source_ip)", + ], +) + + +def _row_to_record(row: sqlite3.Row) -> BottleRecord: + """Build a BottleRecord from a registry row.""" + return BottleRecord( + bottle_id=row["bottle_id"], + source_ip=row["source_ip"], + identity_token=row["identity_token"], + state=row["state"], + created_at=row["created_at"], + metadata=row["metadata"], + ) + + +class RegistryStore(DbStore): + """SQLite-backed registry of live bottles + fail-closed attribution.""" + + def __init__(self, db_path: Path | None = None) -> None: + super().__init__(db_path or default_db_path(), _MIGRATIONS) + + def _connect(self) -> sqlite3.Connection: + """Open a WAL connection (concurrent data-plane reads + writer).""" + conn = super()._connect() + # WAL lets the data-plane attribution reads run concurrently with the + # control-plane writer without blocking; busy_timeout avoids spurious + # "database is locked" under that concurrency (PRD 0070 state tier). + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA busy_timeout=5000") + return conn + + def register( + self, + source_ip: str, + *, + bottle_id: str | None = None, + identity_token: str | None = None, + metadata: str = "", + ) -> BottleRecord: + """Register (or replace) a bottle. Mints a bottle_id / identity token + when not supplied. Live — this is the control-plane reload path.""" + rec = BottleRecord( + bottle_id=bottle_id or secrets.token_hex(8), + source_ip=source_ip, + identity_token=identity_token or new_identity_token(), + state="active", + created_at=time.time(), + metadata=metadata, + ) + with self._connect() as conn: + conn.execute( + "INSERT OR REPLACE INTO orchestrator_bottles " + "(bottle_id, source_ip, identity_token, state, created_at, metadata) " + "VALUES (?, ?, ?, ?, ?, ?)", + ( + rec.bottle_id, + rec.source_ip, + rec.identity_token, + rec.state, + rec.created_at, + rec.metadata, + ), + ) + self._chmod() + return rec + + def deregister(self, bottle_id: str) -> bool: + """Remove a bottle. Returns True if a row was deleted.""" + with self._connect() as conn: + cur = conn.execute( + "DELETE FROM orchestrator_bottles WHERE bottle_id = ?", (bottle_id,) + ) + return cur.rowcount > 0 + + def get(self, bottle_id: str) -> BottleRecord | None: + """Return the bottle by id, or None if absent.""" + with self._connect() as conn: + row = conn.execute( + "SELECT * FROM orchestrator_bottles WHERE bottle_id = ?", (bottle_id,) + ).fetchone() + return _row_to_record(row) if row else None + + def all(self) -> list[BottleRecord]: + """Every registered bottle, oldest first.""" + with self._connect() as conn: + rows = conn.execute( + "SELECT * FROM orchestrator_bottles ORDER BY created_at" + ).fetchall() + return [_row_to_record(r) for r in rows] + + def attribute(self, source_ip: str, identity_token: str) -> BottleRecord | None: + """Fail-closed attribution. Returns the bottle only when exactly one + active record has this source IP AND its identity token matches + (constant-time). Either signal alone is insufficient: an unknown IP, + an ambiguous IP (more than one active bottle — a misconfiguration), + an empty token, or a token mismatch all deny.""" + if not identity_token: + return None + with self._connect() as conn: + rows = conn.execute( + "SELECT * FROM orchestrator_bottles " + "WHERE source_ip = ? AND state = 'active'", + (source_ip,), + ).fetchall() + if len(rows) != 1: + return None + rec = _row_to_record(rows[0]) + if not hmac.compare_digest(rec.identity_token, identity_token): + return None + return rec + + +__all__ = [ + "BottleRecord", + "RegistryStore", + "new_identity_token", + "default_db_path", + "IDENTITY_TOKEN_BYTES", +] diff --git a/tests/unit/test_orchestrator_control_plane.py b/tests/unit/test_orchestrator_control_plane.py new file mode 100644 index 0000000..8310df1 --- /dev/null +++ b/tests/unit/test_orchestrator_control_plane.py @@ -0,0 +1,150 @@ +"""Unit tests for the orchestrator HTTP control plane (PRD 0070). + +Mostly exercises the pure `dispatch()` (socket-free, like the supervise +server tests), plus one real-socket round-trip to prove the handler wiring. +""" + +from __future__ import annotations + +import json +import tempfile +import threading +import unittest +import urllib.request +from pathlib import Path + +from bot_bottle.orchestrator.control_plane import dispatch, make_server +from bot_bottle.orchestrator.registry import RegistryStore + + +def _body(obj: object) -> bytes: + return json.dumps(obj).encode() + + +class TestDispatch(unittest.TestCase): + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.store = RegistryStore(Path(self._tmp.name) / "r.db") + self.store.migrate() + + def tearDown(self) -> None: + self._tmp.cleanup() + + def test_health(self) -> None: + status, payload = dispatch(self.store, "GET", "/health", b"") + self.assertEqual(200, status) + self.assertEqual("ok", payload["status"]) + + def test_register_returns_id_and_token(self) -> None: + status, payload = dispatch( + self.store, "POST", "/bottles", _body({"source_ip": "10.243.0.1"}) + ) + self.assertEqual(201, status) + self.assertTrue(payload["bottle_id"]) + self.assertTrue(payload["identity_token"]) + + def test_register_requires_source_ip(self) -> None: + status, _ = dispatch(self.store, "POST", "/bottles", _body({})) + self.assertEqual(400, status) + + def test_register_rejects_bad_json(self) -> None: + status, _ = dispatch(self.store, "POST", "/bottles", b"{not json") + self.assertEqual(400, status) + + def test_list_redacts_identity_token(self) -> None: + dispatch(self.store, "POST", "/bottles", _body({"source_ip": "10.243.0.1"})) + status, payload = dispatch(self.store, "GET", "/bottles", b"") + self.assertEqual(200, status) + bottles = payload["bottles"] + assert isinstance(bottles, list) + self.assertEqual(1, len(bottles)) + first = bottles[0] + assert isinstance(first, dict) + self.assertNotIn("identity_token", first) + self.assertEqual("10.243.0.1", first["source_ip"]) + + def test_attribute_ok_and_forbidden(self) -> None: + _, reg = dispatch( + self.store, "POST", "/bottles", _body({"source_ip": "10.243.0.5"}) + ) + token = reg["identity_token"] + ok_status, ok = dispatch( + self.store, "POST", "/attribute", + _body({"source_ip": "10.243.0.5", "identity_token": token}), + ) + self.assertEqual(200, ok_status) + self.assertEqual(reg["bottle_id"], ok["bottle_id"]) + + bad_status, _ = dispatch( + self.store, "POST", "/attribute", + _body({"source_ip": "10.243.0.5", "identity_token": "nope"}), + ) + self.assertEqual(403, bad_status) + + def test_attribute_requires_both_fields(self) -> None: + status, _ = dispatch( + self.store, "POST", "/attribute", _body({"source_ip": "10.243.0.5"}) + ) + self.assertEqual(400, status) + + def test_delete(self) -> None: + _, reg = dispatch( + self.store, "POST", "/bottles", _body({"source_ip": "10.243.0.1"}) + ) + bid = reg["bottle_id"] + status, payload = dispatch(self.store, "DELETE", f"/bottles/{bid}", b"") + self.assertEqual(200, status) + self.assertEqual(True, payload["deregistered"]) + self.assertEqual([], self.store.all()) + + def test_delete_missing_404(self) -> None: + status, _ = dispatch(self.store, "DELETE", "/bottles/ghost", b"") + self.assertEqual(404, status) + + def test_unknown_route_404(self) -> None: + status, _ = dispatch(self.store, "GET", "/nope", b"") + self.assertEqual(404, status) + + def test_trailing_slash_normalized(self) -> None: + status, _ = dispatch(self.store, "GET", "/health/", b"") + self.assertEqual(200, status) + + +class TestServerRoundTrip(unittest.TestCase): + def test_http_register_health_attribute(self) -> None: + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + store = RegistryStore(Path(tmp.name) / "r.db") + store.migrate() + server = make_server(store, "127.0.0.1", 0) + self.addCleanup(server.server_close) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + self.addCleanup(server.shutdown) + + host, port = server.server_address[0], server.server_address[1] + base = f"http://{host}:{port}" + + reg = json.load(urllib.request.urlopen( + urllib.request.Request( + f"{base}/bottles", data=_body({"source_ip": "10.243.0.7"}), + method="POST", headers={"Content-Type": "application/json"}, + ), timeout=5, + )) + self.assertTrue(reg["bottle_id"]) + + health = json.load(urllib.request.urlopen(f"{base}/health", timeout=5)) + self.assertEqual("ok", health["status"]) + + attr = json.load(urllib.request.urlopen( + urllib.request.Request( + f"{base}/attribute", + data=_body({"source_ip": "10.243.0.7", "identity_token": reg["identity_token"]}), + method="POST", headers={"Content-Type": "application/json"}, + ), timeout=5, + )) + self.assertEqual(reg["bottle_id"], attr["bottle_id"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_orchestrator_registry.py b/tests/unit/test_orchestrator_registry.py new file mode 100644 index 0000000..6962cb1 --- /dev/null +++ b/tests/unit/test_orchestrator_registry.py @@ -0,0 +1,102 @@ +"""Unit tests for the orchestrator bottle registry + attribution (PRD 0070).""" + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path + +from bot_bottle.orchestrator.registry import ( + BottleRecord, + RegistryStore, + new_identity_token, +) + + +class TestRegistryStore(unittest.TestCase): + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.db = Path(self._tmp.name) / "registry.db" + self.store = RegistryStore(self.db) + self.store.migrate() + + def tearDown(self) -> None: + self._tmp.cleanup() + + def test_register_mints_id_and_token(self) -> None: + rec = self.store.register("10.243.0.1") + self.assertTrue(rec.bottle_id) + self.assertTrue(rec.identity_token) + self.assertEqual("active", rec.state) + self.assertEqual(rec, self.store.get(rec.bottle_id)) + + def test_identity_tokens_are_unique(self) -> None: + tokens = {new_identity_token() for _ in range(200)} + self.assertEqual(200, len(tokens)) + a = self.store.register("10.243.0.1") + b = self.store.register("10.243.0.3") + self.assertNotEqual(a.identity_token, b.identity_token) + + def test_all_and_deregister(self) -> None: + a = self.store.register("10.243.0.1") + b = self.store.register("10.243.0.3") + self.assertEqual({a.bottle_id, b.bottle_id}, {r.bottle_id for r in self.store.all()}) + self.assertTrue(self.store.deregister(a.bottle_id)) + self.assertEqual([b.bottle_id], [r.bottle_id for r in self.store.all()]) + + def test_deregister_missing_is_false(self) -> None: + self.assertFalse(self.store.deregister("nope")) + + def test_explicit_id_replaces(self) -> None: + self.store.register("10.243.0.1", bottle_id="fixed", metadata="a") + self.store.register("10.243.0.9", bottle_id="fixed", metadata="b") + rec = self.store.get("fixed") + assert rec is not None + self.assertEqual("10.243.0.9", rec.source_ip) + self.assertEqual("b", rec.metadata) + self.assertEqual(1, len(self.store.all())) + + def test_attribute_success_requires_ip_and_token(self) -> None: + rec = self.store.register("10.243.0.1") + got = self.store.attribute("10.243.0.1", rec.identity_token) + assert got is not None + self.assertEqual(rec.bottle_id, got.bottle_id) + + def test_attribute_wrong_token_denied(self) -> None: + self.store.register("10.243.0.1") + self.assertIsNone(self.store.attribute("10.243.0.1", "wrong-token")) + + def test_attribute_unknown_ip_denied(self) -> None: + rec = self.store.register("10.243.0.1") + self.assertIsNone(self.store.attribute("10.243.9.9", rec.identity_token)) + + def test_attribute_empty_token_denied(self) -> None: + self.store.register("10.243.0.1") + self.assertIsNone(self.store.attribute("10.243.0.1", "")) + + def test_attribute_ambiguous_ip_denied(self) -> None: + # Two active bottles on one source IP is a misconfiguration — deny + # rather than guess (fail-closed), even with a valid token. + a = self.store.register("10.243.0.1", bottle_id="a") + self.store.register("10.243.0.1", bottle_id="b") + self.assertIsNone(self.store.attribute("10.243.0.1", a.identity_token)) + + def test_state_persists_across_reopen(self) -> None: + rec = self.store.register("10.243.0.1") + reopened = RegistryStore(self.db) + got = reopened.get(rec.bottle_id) + assert got is not None + self.assertEqual(rec, got) + # Attribution works against the reopened (durable) store too. + self.assertIsNotNone(reopened.attribute("10.243.0.1", rec.identity_token)) + + def test_redacted_hides_token(self) -> None: + rec = BottleRecord( + bottle_id="x", source_ip="10.243.0.1", identity_token="secret" + ) + self.assertNotIn("identity_token", rec.redacted()) + self.assertEqual("10.243.0.1", rec.redacted()["source_ip"]) + + +if __name__ == "__main__": + unittest.main()