bot-bottle/tests/unit/test_orchestrator_control_plane.py

"""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 secrets
import tempfile
import threading
import unittest
import urllib.error
import urllib.request
from pathlib import Path
from unittest.mock import patch

from bot_bottle.orchestrator.broker import StubBroker
from bot_bottle.orchestrator.control_plane import dispatch, make_server
from bot_bottle.orchestrator.registry import RegistryStore
from bot_bottle.orchestrator.service import Orchestrator
from bot_bottle.store_manager import StoreManager
from bot_bottle.supervise import (
    Proposal,
    TOOL_EGRESS_ALLOW,
    sha256_hex,
    write_proposal,
)


def _body(obj: object) -> bytes:
    return json.dumps(obj).encode()


def _orchestrator(db_path: Path) -> Orchestrator:
    store = RegistryStore(db_path)
    store.migrate()
    secret = secrets.token_bytes(16)
    return Orchestrator(store, StubBroker(secret), secret)


class TestDispatch(unittest.TestCase):
    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.orch = _orchestrator(Path(self._tmp.name) / "r.db")

    def tearDown(self) -> None:
        self._tmp.cleanup()

    def test_health(self) -> None:
        status, payload = dispatch(self.orch, "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.orch, "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.orch, "POST", "/bottles", _body({}))
        self.assertEqual(400, status)

    def test_register_rejects_bad_json(self) -> None:
        status, _ = dispatch(self.orch, "POST", "/bottles", b"{not json")
        self.assertEqual(400, status)

    def test_list_redacts_identity_token(self) -> None:
        dispatch(self.orch, "POST", "/bottles", _body({"source_ip": "10.243.0.1"}))
        status, payload = dispatch(self.orch, "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.orch, "POST", "/bottles", _body({"source_ip": "10.243.0.5"})
        )
        token = reg["identity_token"]
        ok_status, ok = dispatch(
            self.orch, "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.orch, "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.orch, "POST", "/attribute", _body({"source_ip": "10.243.0.5"})
        )
        self.assertEqual(400, status)

    def test_delete(self) -> None:
        _, reg = dispatch(
            self.orch, "POST", "/bottles", _body({"source_ip": "10.243.0.1"})
        )
        bid = reg["bottle_id"]
        status, payload = dispatch(self.orch, "DELETE", f"/bottles/{bid}", b"")
        self.assertEqual(200, status)
        self.assertEqual(True, payload["torn_down"])
        self.assertEqual([], self.orch.registry.all())

    def test_delete_missing_404(self) -> None:
        status, _ = dispatch(self.orch, "DELETE", "/bottles/ghost", b"")
        self.assertEqual(404, status)

    def test_unknown_route_404(self) -> None:
        status, _ = dispatch(self.orch, "GET", "/nope", b"")
        self.assertEqual(404, status)

    def test_trailing_slash_normalized(self) -> None:
        status, _ = dispatch(self.orch, "GET", "/health/", b"")
        self.assertEqual(200, status)

    def test_gateway_status_unconfigured(self) -> None:
        status, payload = dispatch(self.orch, "GET", "/gateway", b"")
        self.assertEqual(200, status)
        self.assertEqual(False, payload["configured"])

    def test_launch_stores_policy_and_resolve_returns_it(self) -> None:
        _, reg = dispatch(
            self.orch, "POST", "/bottles",
            _body({"source_ip": "10.243.0.5", "policy": '{"a":1}'}),
        )
        status, payload = dispatch(
            self.orch, "POST", "/resolve",
            _body({"source_ip": "10.243.0.5", "identity_token": reg["identity_token"]}),
        )
        self.assertEqual(200, status)
        self.assertEqual(reg["bottle_id"], payload["bottle_id"])
        self.assertEqual('{"a":1}', payload["policy"])

    def test_resolve_forbidden_on_bad_token(self) -> None:
        dispatch(self.orch, "POST", "/bottles", _body({"source_ip": "10.243.0.5"}))
        status, _ = dispatch(
            self.orch, "POST", "/resolve",
            _body({"source_ip": "10.243.0.5", "identity_token": "nope"}),
        )
        self.assertEqual(403, status)

    def test_put_policy_updates_live(self) -> None:
        _, reg = dispatch(self.orch, "POST", "/bottles", _body({"source_ip": "10.243.0.1"}))
        bid = reg["bottle_id"]
        status, payload = dispatch(
            self.orch, "PUT", f"/bottles/{bid}/policy", _body({"policy": '{"v":9}'})
        )
        self.assertEqual(200, status)
        self.assertEqual(True, payload["updated"])
        _, resolved = dispatch(
            self.orch, "POST", "/resolve",
            _body({"source_ip": "10.243.0.1", "identity_token": reg["identity_token"]}),
        )
        self.assertEqual('{"v":9}', resolved["policy"])

    def test_put_policy_missing_bottle_404(self) -> None:
        status, _ = dispatch(
            self.orch, "PUT", "/bottles/ghost/policy", _body({"policy": "{}"})
        )
        self.assertEqual(404, status)

    def test_put_policy_requires_string(self) -> None:
        _, reg = dispatch(self.orch, "POST", "/bottles", _body({"source_ip": "10.243.0.1"}))
        status, _ = dispatch(
            self.orch, "PUT", f"/bottles/{reg['bottle_id']}/policy", _body({})
        )
        self.assertEqual(400, status)

    def test_resolve_without_token_denies(self) -> None:
        # Mandatory token: source-IP alone no longer resolves (fail-closed 403).
        dispatch(
            self.orch, "POST", "/bottles",
            _body({"source_ip": "10.243.0.5", "policy": "P"}),
        )
        status, _ = dispatch(
            self.orch, "POST", "/resolve", _body({"source_ip": "10.243.0.5"})
        )
        self.assertEqual(403, status)

    def test_resolve_with_matching_token(self) -> None:
        _, reg = dispatch(
            self.orch, "POST", "/bottles",
            _body({"source_ip": "10.243.0.6", "policy": "P"}),
        )
        status, payload = dispatch(
            self.orch, "POST", "/resolve",
            _body({"source_ip": "10.243.0.6", "identity_token": reg["identity_token"]}),
        )
        self.assertEqual(200, status)
        self.assertEqual(reg["bottle_id"], payload["bottle_id"])
        self.assertEqual("P", payload["policy"])

    def test_resolve_requires_source_ip(self) -> None:
        status, _ = dispatch(self.orch, "POST", "/resolve", _body({}))
        self.assertEqual(400, status)


class TestServerRoundTrip(unittest.TestCase):
    def test_http_register_health_attribute(self) -> None:
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        orch = _orchestrator(Path(tmp.name) / "r.db")
        server = make_server(orch, "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"])


class TestControlPlaneAuth(unittest.TestCase):
    """The per-host control-plane secret (issue #400): every route but /health
    is a trusted-caller op an agent must not be able to drive just because it
    can reach the port."""

    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self._tmp.cleanup)
        self.orch = _orchestrator(Path(self._tmp.name) / "r.db")

    def test_health_is_public_even_unauthorized(self) -> None:
        status, _ = dispatch(self.orch, "GET", "/health", b"", authorized=False)
        self.assertEqual(200, status)

    def test_unauthorized_denies_every_other_route(self) -> None:
        for method, path, body in [
            ("GET", "/bottles", b""),
            ("POST", "/bottles", _body({"source_ip": "10.0.0.1"})),
            ("PUT", "/bottles/x/policy", _body({"policy": "routes: []"})),
            ("DELETE", "/bottles/x", b""),
            ("POST", "/resolve", _body({"source_ip": "10.0.0.1", "identity_token": "t"})),
            ("POST", "/attribute", _body({"source_ip": "10.0.0.1", "identity_token": "t"})),
            ("GET", "/supervise/proposals", b""),
            ("POST", "/supervise/respond", _body({"proposal_id": "p", "bottle_slug": "s", "decision": "approve"})),
        ]:
            status, _ = dispatch(self.orch, method, path, body, authorized=False)
            self.assertEqual(401, status, f"{method} {path} should be 401 unauthorized")

    def test_deny_happens_before_the_registry_is_touched(self) -> None:
        """An unauthorized DELETE must not tear a bottle down. 401, and the
        bottle is still there."""
        rec = self.orch.registry.register("10.0.0.9", policy="", metadata="")
        status, _ = dispatch(
            self.orch, "DELETE", f"/bottles/{rec.bottle_id}", b"", authorized=False)
        self.assertEqual(401, status)
        self.assertIsNotNone(self.orch.registry.get(rec.bottle_id))

    def _server_with_secret(self, secret: str):
        with patch.dict("os.environ", {"BOT_BOTTLE_CONTROL_PLANE_TOKEN": secret}):
            server = make_server(self.orch, "127.0.0.1", 0)
        self.addCleanup(server.server_close)
        threading.Thread(target=server.serve_forever, daemon=True).start()
        self.addCleanup(server.shutdown)
        host, port = server.server_address[0], server.server_address[1]
        return f"http://{host}:{port}"

    def _status(self, url: str, *, header: str | None = None) -> int:
        req = urllib.request.Request(url)
        if header is not None:
            req.add_header("x-bot-bottle-control-auth", header)
        try:
            return urllib.request.urlopen(req, timeout=5).status
        except urllib.error.HTTPError as e:
            return e.code

    def test_configured_server_enforces_the_header_over_http(self) -> None:
        base = self._server_with_secret("s3cret-admin")
        # /health is public — no header needed.
        self.assertEqual(200, self._status(f"{base}/health"))
        # /bottles requires the secret.
        self.assertEqual(401, self._status(f"{base}/bottles"))
        self.assertEqual(401, self._status(f"{base}/bottles", header="wrong"))
        self.assertEqual(200, self._status(f"{base}/bottles", header="s3cret-admin"))

    def test_unconfigured_server_runs_open(self) -> None:
        """No secret set (tests / nft-protected Firecracker): open mode, so the
        existing round-trip and unit behavior are unchanged."""
        with patch.dict("os.environ", {}, clear=False):
            import os
            os.environ.pop("BOT_BOTTLE_CONTROL_PLANE_TOKEN", None)
            server = make_server(self.orch, "127.0.0.1", 0)
        self.addCleanup(server.server_close)
        self.assertTrue(server.is_authorized(""))
        self.assertTrue(server.is_authorized("anything"))


class TestDispatchSupervise(unittest.TestCase):
    """The /supervise/* routes over the pure dispatch()."""

    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        root = Path(self._tmp.name)
        db = root / "db" / "bot-bottle.db"
        db.parent.mkdir(parents=True)
        self._env = patch.dict("os.environ", {
            "BOT_BOTTLE_ROOT": str(root),
            "SUPERVISE_DB_PATH": str(db),
        })
        self._env.start()
        self.store = RegistryStore(db)
        self.store.migrate()
        StoreManager(db).migrate()
        secret = secrets.token_bytes(16)
        self.orch = Orchestrator(self.store, StubBroker(secret), secret)

    def tearDown(self) -> None:
        self._env.stop()
        self._tmp.cleanup()

    def _queue(self, slug: str, proposed: str) -> str:
        self.store.register(
            "10.243.0.1", metadata=json.dumps({"slug": slug}), policy="routes: []\n")
        p = Proposal.new(
            bottle_slug=slug, tool=TOOL_EGRESS_ALLOW, proposed_file=proposed,
            justification="need it", current_file_hash=sha256_hex(proposed))
        write_proposal(p)
        return p.id

    def test_list_pending(self) -> None:
        pid = self._queue("demo", "routes:\n  - host: google.com\n")
        status, payload = dispatch(self.orch, "GET", "/supervise/proposals", b"")
        self.assertEqual(200, status)
        proposals = payload["proposals"]
        assert isinstance(proposals, list)
        self.assertEqual(pid, proposals[0]["id"])

    def test_respond_approve_applies_and_clears(self) -> None:
        pid = self._queue("demo", "routes:\n  - host: google.com\n")
        status, payload = dispatch(
            self.orch, "POST", "/supervise/respond",
            _body({"proposal_id": pid, "bottle_slug": "demo", "decision": "approve"}),
        )
        self.assertEqual(200, status)
        self.assertTrue(payload["responded"])
        _, listing = dispatch(self.orch, "GET", "/supervise/proposals", b"")
        self.assertEqual([], listing["proposals"])

    def test_respond_requires_fields(self) -> None:
        status, _ = dispatch(
            self.orch, "POST", "/supervise/respond", _body({"decision": "approve"}))
        self.assertEqual(400, status)

    def test_respond_unknown_proposal_conflicts(self) -> None:
        status, payload = dispatch(
            self.orch, "POST", "/supervise/respond",
            _body({"proposal_id": "ghost", "bottle_slug": "demo", "decision": "approve"}),
        )
        self.assertEqual(409, status)
        self.assertIn("no such proposal", str(payload["error"]))


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