broker/tests/test_client.py

import ast
import base64
import contextlib
import datetime
import io
import json
import os
import sys
import tempfile
import time
import unittest
import urllib.error
import urllib.request
import uuid as uuidlib

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)))

import coreci_broker
import coreci_broker_client

ADMIN_USER = "admin"
ADMIN_PASS = "correct-horse-battery-staple"
APPS_PATH = "/api/broker/app_registrations"
USERS_PATH = "/api/broker/user_allowlist"
PROXMOX_SECRET = {"url": "https://pve.coreci.dev:8006", "token": "pve-tok"}
GITEA_SECRET = {"url": "https://git.coreci.dev", "token": "git-tok"}
CLIENT_SOURCE_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), os.pardir, "coreci_broker_client.py")
)
ALLOWED_CLIENT_IMPORTS = {"urllib.request", "urllib.parse", "json", "time"}
SECRETS_PROXMOX_PATH = "/api/broker/secrets/coreci/proxmox"


def _basic_header(username, password):
    raw = "%s:%s" % (username, password)
    return "Basic " + base64.b64encode(raw.encode("utf-8")).decode("ascii")


class ClientHygieneTests(unittest.TestCase):
    def setUp(self):
        with open(CLIENT_SOURCE_PATH, "r", encoding="utf-8") as handle:
            self.source = handle.read()
        self.tree = ast.parse(self.source)

    def test_imports_are_stdlib_only(self):
        modules = set()
        for node in ast.walk(self.tree):
            if isinstance(node, ast.Import):
                modules.update(alias.name for alias in node.names)
            elif isinstance(node, ast.ImportFrom):
                modules.add(node.module or "")
        self.assertIn("urllib.request", modules)
        for module in modules:
            self.assertIn(module, ALLOWED_CLIENT_IMPORTS, module)

    def test_single_module_docstring_and_no_comments(self):
        docstring = ast.get_docstring(self.tree)
        self.assertIsNotNone(docstring)
        self.assertNotIn("\n", docstring)
        for node in ast.walk(self.tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                self.assertIsNone(ast.get_docstring(node))
        for number, line in enumerate(self.source.splitlines(), 1):
            self.assertFalse(line.lstrip().startswith("#"), "comment on line %d" % number)

    def test_broker_error_carries_status_and_message(self):
        error = coreci_broker_client.BrokerError(403, "forbidden")
        self.assertEqual(403, error.status)
        self.assertEqual("forbidden", error.message)
        self.assertEqual("forbidden", str(error))


class _ClientCase(unittest.TestCase):
    extra_env = {}

    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.tmp = tmp.name
        self.creds_path = os.path.join(self.tmp, "creds.json")
        with open(self.creds_path, "w", encoding="utf-8") as handle:
            handle.write(json.dumps(
                {"coreci/proxmox": PROXMOX_SECRET, "coreci/gitea": GITEA_SECRET}
            ))
        env = {
            "BROKER_ENV_PATH": os.path.join(self.tmp, ".env"),
            "BROKER_DB_PATH": os.path.join(self.tmp, "broker.db"),
            "BROKER_AUDIT_LOG_PATH": os.path.join(self.tmp, "broker.audit.log"),
            "BROKER_ADMIN_USER": ADMIN_USER,
            "BROKER_ADMIN_PASS": ADMIN_PASS,
            "BROKER_CREDS_SOURCE": "file",
            "BROKER_CREDS_FILE": self.creds_path,
        }
        env.update(self.extra_env)
        with contextlib.redirect_stdout(io.StringIO()):
            self.broker = coreci_broker.boot(env)
        self.addCleanup(self.broker.audit.close)
        self.addCleanup(self.broker.store.close)
        self.server = coreci_broker.serve("127.0.0.1", 0, self.broker)
        self.addCleanup(self.server.shutdown)
        self.addCleanup(self.server.server_close)
        self.base = "http://127.0.0.1:%d" % self.server.server_address[1]
        self.admin = {"Authorization": _basic_header(ADMIN_USER, ADMIN_PASS)}
        self.app_uuid = str(uuidlib.uuid4())
        self.user_id = "ceo"
        self._seed()

    def _raw_request(self, method, path, payload=None, headers=None):
        data = None if payload is None else json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(self.base + path, data=data, method=method)
        for key, value in (headers or {}).items():
            request.add_header(key, value)
        try:
            with urllib.request.urlopen(request, timeout=10) as response:
                raw = response.read().decode("utf-8")
                return response.status, json.loads(raw) if raw else None
        except urllib.error.HTTPError as error:
            raw = error.read().decode("utf-8")
            return error.code, json.loads(raw) if raw else None

    def _seed(self):
        status, _ = self._raw_request("POST", APPS_PATH, {"uuid": self.app_uuid}, self.admin)
        self.assertEqual(201, status)
        status, _ = self._raw_request("POST", USERS_PATH, {"user_id": self.user_id}, self.admin)
        self.assertEqual(201, status)

    def _client(self, agent_id="proxmox", tags=("coreci",)):
        return coreci_broker_client.BrokerClient(
            self.base, agent_id, self.app_uuid, self.user_id, list(tags)
        )

    def _count_requests(self, client):
        calls = []
        original = client._request

        def counting(method, path, payload=None, headers=None):
            calls.append((method, path))
            return original(method, path, payload, headers)

        def restore():
            client._request = original

        client._request = counting
        return calls, restore


class ClientFlowTests(_ClientCase):
    def test_ensure_registered_from_empty_state_populates_token_fields(self):
        client = self._client()
        state = {}
        returned = client.ensure_registered(state)
        self.assertIs(state, returned)
        token = state["current_token"]
        self.assertEqual(44, len(token))
        self.assertTrue(token.isascii())
        expires = datetime.datetime.fromisoformat(state["token_expires_at"])
        self.assertGreater(expires, datetime.datetime.now(datetime.timezone.utc))
        self.assertEqual("proxmox", state["agent_id"])
        self.assertEqual(self.app_uuid, state["app_registration_uuid"])
        self.assertEqual(self.user_id, state["user_id"])
        self.assertEqual(["coreci"], state["tags"])
        status, body = self._raw_request(
            "GET", SECRETS_PROXMOX_PATH,
            headers={"Authorization": "App %s %s" % (self.app_uuid, token)},
        )
        self.assertEqual(200, status)
        self.assertEqual(PROXMOX_SECRET, body["data"])

    def test_first_register_merges_pepper_into_state(self):
        client = self._client()
        state = client.ensure_registered({})
        self.assertIn("pepper", state)
        self.assertEqual(
            self.broker.store.get_config("BROKER_INSTANCE_PEPPER"), state["pepper"]
        )

    def test_ensure_registered_with_live_token_makes_no_network_call(self):
        client = self._client()
        state = client.ensure_registered({})
        snapshot = dict(state)
        calls, restore = self._count_requests(client)
        try:
            returned = client.ensure_registered(state)
        finally:
            restore()
        self.assertIs(state, returned)
        self.assertEqual([], calls)
        self.assertEqual(snapshot, state)

    def test_reregister_with_valid_token_pepperless_response_does_not_crash(self):
        client = self._client()
        state = client.ensure_registered({})
        pepper, token, expires = state["pepper"], state["current_token"], state["token_expires_at"]
        status, body = self._raw_request(
            "POST", coreci_broker_client.REGISTER_PATH,
            {"agent_id": "proxmox", "app_registration_uuid": self.app_uuid,
             "user_id": self.user_id, "tags": ["coreci"]},
            {"Authorization": "App %s %s" % (self.app_uuid, token)},
        )
        self.assertEqual(200, status)
        self.assertNotIn("pepper", body)
        returned = client.register(state)
        self.assertIs(state, returned)
        self.assertEqual(pepper, state["pepper"])
        self.assertEqual(token, state["current_token"])
        self.assertEqual(expires, state["token_expires_at"])

    def test_put_then_get_state_roundtrip_through_broker(self):
        client = self._client()
        state = client.ensure_registered({})
        state["scratchpad"] = {"calls": 3, "notes": ["x"]}
        self.assertTrue(client.put_state(state))
        document = client.get_state(state)
        self.assertEqual(state, document)

    def test_get_state_returns_none_before_first_put(self):
        client = self._client()
        state = client.ensure_registered({})
        self.assertIsNone(client.get_state(state))

    def test_read_secret_returns_secret_document(self):
        client = self._client()
        state = client.ensure_registered({})
        document = client.read_secret("coreci/proxmox", state)
        self.assertEqual("coreci/proxmox", document["path"])
        self.assertEqual(PROXMOX_SECRET, document["data"])
        self.assertRegex(document["version"], r"^[0-9a-f]{8}$")
        self.assertIn("updated_at", document)

    def test_cross_path_read_secret_raises_403_broker_error(self):
        client = self._client()
        state = client.ensure_registered({})
        with self.assertRaises(coreci_broker_client.BrokerError) as caught:
            client.read_secret("coreci/gitea", state)
        self.assertEqual(403, caught.exception.status)
        self.assertTrue(caught.exception.message)

    def test_missing_secret_read_raises_404_broker_error(self):
        client = self._client("nonexistent")
        state = client.ensure_registered({})
        with self.assertRaises(coreci_broker_client.BrokerError) as caught:
            client.read_secret("coreci/nonexistent", state)
        self.assertEqual(404, caught.exception.status)

    def test_read_secret_from_cold_state_registers_then_retries_once(self):
        client = self._client()
        state = {}
        calls, restore = self._count_requests(client)
        try:
            document = client.read_secret("coreci/proxmox", state)
        finally:
            restore()
        self.assertEqual(PROXMOX_SECRET, document["data"])
        self.assertEqual(
            [("GET", SECRETS_PROXMOX_PATH),
             ("POST", coreci_broker_client.REGISTER_PATH),
             ("GET", SECRETS_PROXMOX_PATH)],
            calls,
        )
        self.assertIn("current_token", state)

    def test_get_state_with_foreign_agent_token_is_403_identity_enforced(self):
        intruder = self._client("proxmox")
        other = self._client("gitea-agent")
        other_state = other.ensure_registered({})
        intruder.ensure_registered({})
        with self.assertRaises(coreci_broker_client.BrokerError) as caught:
            intruder.get_state({"current_token": other_state["current_token"]})
        self.assertEqual(403, caught.exception.status)


class TtlAutoReregisterTests(_ClientCase):
    extra_env = {"BROKER_TOKEN_TTL": "1"}

    def test_expired_token_secret_read_auto_reregisters_and_retries(self):
        client = self._client()
        state = {}
        first = client.read_secret("coreci/proxmox", state)
        self.assertEqual(PROXMOX_SECRET, first["data"])
        first_expires = datetime.datetime.fromisoformat(state["token_expires_at"])
        time.sleep(1.2)
        second = client.read_secret("coreci/proxmox", state)
        self.assertEqual(PROXMOX_SECRET, second["data"])
        second_expires = datetime.datetime.fromisoformat(state["token_expires_at"])
        self.assertGreater(second_expires, first_expires)
        self.assertIn("pepper", state)


class StaticDeliveryTests(_ClientCase):
    def test_fetch_client_module_returns_own_source(self):
        with open(CLIENT_SOURCE_PATH, "r", encoding="utf-8") as handle:
            source = handle.read()
        fetched = coreci_broker_client.fetch_client_module(self.base)
        self.assertEqual(source, fetched)
        self.assertIn("class BrokerClient", fetched)

    def test_client_module_served_unauthenticated(self):
        with open(CLIENT_SOURCE_PATH, "r", encoding="utf-8") as handle:
            source = handle.read()
        with urllib.request.urlopen(
            self.base + "/static/coreci_broker_client.py", timeout=10
        ) as response:
            self.assertEqual(200, response.status)
            self.assertEqual(source, response.read().decode("utf-8"))


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