broker/tests/test_register.py

import base64
import contextlib
import datetime
import io
import json
import os
import sys
import tempfile
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

ADMIN_USER = "admin"
ADMIN_PASS = "correct-horse-battery-staple"
REGISTER_PATH = "/api/broker/agents/register"
APPS_PATH = "/api/broker/app_registrations"
USERS_PATH = "/api/broker/user_allowlist"
WHOAMI_PATH = "/api/broker/auth/whoami"
DEFAULT_TTL = 86400


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


def _seed_direct(broker, app_uuid, user_id):
    broker.store.add_row("app_registrations", app_uuid, None)
    broker.store.add_row("user_allowlist", user_id, None)


class AgentRegisterTests(unittest.TestCase):
    def setUp(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        self.tmp = tmp.name
        self.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,
        }
        with contextlib.redirect_stdout(io.StringIO()):
            self.broker = coreci_broker.boot(dict(self.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(self.app_uuid, self.user_id)

    def _request(self, method, path, payload=None, headers=None, base=None):
        data = None if payload is None else json.dumps(payload).encode("utf-8")
        request = urllib.request.Request(
            (base or 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:
                return response.status, json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as error:
            return error.code, json.loads(error.read().decode("utf-8"))

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

    def _register(self, agent_id="agent-1", tags=("coreci",), app_uuid=None,
                  user_id=None, headers=None, payload=None):
        if payload is None:
            payload = {
                "agent_id": agent_id,
                "app_registration_uuid": app_uuid or self.app_uuid,
                "user_id": user_id or self.user_id,
                "tags": list(tags),
            }
        return self._request("POST", REGISTER_PATH, payload, headers)

    def _agent_row(self, agent_id="agent-1", app_uuid=None):
        return self.broker.store.get_agent(agent_id, app_uuid or self.app_uuid)

    def _audit(self, action=None, result=None):
        with open(self.env["BROKER_AUDIT_LOG_PATH"], "r", encoding="utf-8") as handle:
            entries = [json.loads(raw) for raw in handle if raw.strip()]
        return [
            entry for entry in entries
            if (action is None or entry["action"] == action)
            and (result is None or entry["result"] == result)
        ]

    def _pepper(self):
        return self.broker.store.get_config("BROKER_INSTANCE_PEPPER")

    def _app_header(self, token, app_uuid=None):
        return {"Authorization": "App %s %s" % (app_uuid or self.app_uuid, token)}

    def _ttl_delta_seconds(self, body):
        expires = datetime.datetime.fromisoformat(body["token_expires_at"])
        now = datetime.datetime.now(datetime.timezone.utc)
        return (expires - now).total_seconds()

    def test_gate_matrix_unknown_and_revoked_uuid_and_user(self):
        unknown_uuid = str(uuidlib.uuid4())
        status, body = self._register(app_uuid=unknown_uuid)
        self.assertEqual(403, status)
        self.assertEqual(403, body["error"])
        self.assertTrue(body["message"])
        self._request("DELETE", APPS_PATH + "/" + self.app_uuid, headers=self.admin)
        status, body = self._register()
        self.assertEqual(403, status)
        self._request("POST", APPS_PATH + "/" + self.app_uuid + "/reinstate", headers=self.admin)
        status, body = self._register(user_id="stranger")
        self.assertEqual(403, status)
        self._request("DELETE", USERS_PATH + "/" + self.user_id, headers=self.admin)
        status, body = self._register()
        self.assertEqual(403, status)
        denied = self._audit("register", "denied")
        self.assertEqual(4, len(denied))
        for entry in denied:
            self.assertEqual(REGISTER_PATH, entry["path"])
            self.assertEqual("agent-1", entry["agent_id"])
        self.assertEqual(unknown_uuid, denied[0]["app_registration_uuid"])
        self.assertEqual(self.user_id, denied[0]["user_id"])
        self.assertEqual(self.app_uuid, denied[1]["app_registration_uuid"])
        self.assertIsNone(self._agent_row())
        self._request("POST", USERS_PATH + "/" + self.user_id + "/reinstate", headers=self.admin)
        status, body = self._register()
        self.assertEqual(201, status)
        self.assertTrue(body["created"])
        self.assertIsNotNone(self._agent_row())

    def test_malformed_bodies_are_400(self):
        base = {"agent_id": "a", "app_registration_uuid": self.app_uuid,
                "user_id": self.user_id, "tags": ["x"]}
        variants = (
            {}, {"agent_id": "a"}, {"agent_id": "a", "app_registration_uuid": self.app_uuid},
            dict(base, agent_id=None), dict(base, agent_id=123), dict(base, agent_id=""),
            dict(base, app_registration_uuid=None), dict(base, app_registration_uuid=[]),
            dict(base, user_id=7), dict(base, user_id=""), dict(base, tags=None),
            dict(base, tags="x"), dict(base, tags=["a", 1]), dict(base, tags=[["nested"]]),
            dict(base, tags={"a": 1}), [1, 2, 3], "nope",
        )
        for payload in variants:
            status, body = self._request("POST", REGISTER_PATH, payload)
            self.assertEqual(400, status, repr(payload))
            self.assertEqual(400, body["error"])
            self.assertTrue(body["message"])
        self.assertIsNone(self._agent_row())
        self.assertEqual([], self._audit("register"))

    def test_first_register_creates_row_pepper_and_default_ttl(self):
        status, body = self._register(agent_id="px-1", tags=("proxmox", "readonly"))
        self.assertEqual(201, status)
        self.assertTrue(body["created"])
        self.assertEqual(self._pepper(), body["pepper"])
        expected = coreci_broker.derive_token(
            self._pepper(), self.app_uuid, self.user_id, "px-1", ["proxmox", "readonly"]
        )
        self.assertEqual(expected, body["token"])
        self.assertEqual(44, len(body["token"]))
        self.assertTrue(body["token"].isascii())
        delta = self._ttl_delta_seconds(body)
        self.assertGreater(delta, DEFAULT_TTL - 300)
        self.assertLess(delta, DEFAULT_TTL + 300)
        row = self._agent_row("px-1")
        self.assertNotEqual(body["token"], row[4])
        self.assertTrue(coreci_broker.verify_token(body["token"], row[4]))
        self.assertEqual(self.user_id, row[2])
        self.assertEqual(["proxmox", "readonly"], json.loads(row[3]))
        self.assertEqual(body["token_expires_at"], row[5])
        self.assertTrue(row[6])
        self.assertTrue(row[7])
        entries = self._audit("register", "ok")
        self.assertEqual(1, len(entries))
        entry = entries[0]
        self.assertEqual(self.user_id, entry["user_id"])
        self.assertEqual(self.app_uuid, entry["app_registration_uuid"])
        self.assertEqual("px-1", entry["agent_id"])
        self.assertEqual(REGISTER_PATH, entry["path"])

    def test_reregister_unchanged_reissues_same_token_without_rehash(self):
        _, first = self._register()
        stored_hash = self._agent_row()[4]
        seen_at = self._agent_row()[7]
        status, second = self._register()
        self.assertEqual(200, status)
        self.assertFalse(second["created"])
        self.assertEqual(first["token"], second["token"])
        self.assertEqual(first["token_expires_at"], second["token_expires_at"])
        row = self._agent_row()
        self.assertEqual(stored_hash, row[4])
        self.assertGreaterEqual(row[7], seen_at)
        self.assertTrue(coreci_broker.verify_token(first["token"], row[4]))
        self.assertEqual(2, len(self._audit("register", "ok")))

    def test_reregister_pepper_rule_tokenless_vs_valid_app_token(self):
        _, first = self._register()
        status, tokenless = self._register()
        self.assertEqual(200, status)
        self.assertEqual(first["token"], tokenless["token"])
        self.assertIn("pepper", tokenless)
        status, with_token = self._register(headers=self._app_header(first["token"]))
        self.assertEqual(200, status)
        self.assertNotIn("pepper", with_token)
        self.assertEqual(first["token"], with_token["token"])
        status, bad_token = self._register(
            headers={"Authorization": "App %s wrong-token" % self.app_uuid}
        )
        self.assertEqual(200, status)
        self.assertIn("pepper", bad_token)

    def test_changed_tags_overwrite_hash_token_and_expiry(self):
        _, first = self._register(tags=("red",))
        status, second = self._register(tags=("red", "blue"))
        self.assertEqual(200, status)
        self.assertFalse(second["created"])
        self.assertNotEqual(first["token"], second["token"])
        row = self._agent_row()
        self.assertFalse(coreci_broker.verify_token(first["token"], row[4]))
        self.assertTrue(coreci_broker.verify_token(second["token"], row[4]))
        self.assertEqual(["blue", "red"], json.loads(row[3]))
        self.assertEqual(second["token_expires_at"], row[5])

    def test_changed_user_overwrites_row_identity(self):
        self._request("POST", USERS_PATH, {"user_id": "coo"}, self.admin)
        _, first = self._register()
        status, second = self._register(user_id="coo")
        self.assertEqual(200, status)
        self.assertNotEqual(first["token"], second["token"])
        row = self._agent_row()
        self.assertEqual("coo", row[2])
        self.assertFalse(coreci_broker.verify_token(first["token"], row[4]))
        self.assertTrue(coreci_broker.verify_token(second["token"], row[4]))

    def test_expired_row_reregister_rehashes_and_extends(self):
        _, first = self._register()
        row = self._agent_row()
        old_hash, first_registered = row[4], row[6]
        past = (datetime.datetime.now(datetime.timezone.utc)
                - datetime.timedelta(seconds=60)).isoformat()
        self.broker.store.conn.execute(
            "UPDATE agent_registrations SET token_expires_at = ? WHERE agent_id = ?",
            (past, "agent-1"),
        )
        self.broker.store.conn.commit()
        status, second = self._register()
        self.assertEqual(200, status)
        self.assertFalse(second["created"])
        self.assertEqual(first["token"], second["token"])
        delta = self._ttl_delta_seconds(second)
        self.assertGreater(delta, DEFAULT_TTL - 300)
        row = self._agent_row()
        self.assertNotEqual(old_hash, row[4])
        self.assertEqual(first_registered, row[6])
        self.assertTrue(coreci_broker.verify_token(first["token"], row[4]))
        self.assertEqual(second["token_expires_at"], row[5])
        self.assertIn("pepper", second)

    def test_revoked_app_blocks_new_registrations_rows_untouched(self):
        _, first = self._register()
        row_before = self._agent_row()
        self._request("DELETE", APPS_PATH + "/" + self.app_uuid, headers=self.admin)
        status, body = self._register(agent_id="fresh-agent")
        self.assertEqual(403, status)
        status, body = self._register()
        self.assertEqual(403, status)
        self.assertEqual(2, len(self._audit("register", "denied")))
        row_after = self._agent_row()
        self.assertEqual(row_before[4], row_after[4])
        self.assertEqual(row_before[5], row_after[5])
        self.assertTrue(coreci_broker.verify_token(first["token"], row_after[4]))

    def test_token_ttl_override_honored(self):
        env = dict(self.env)
        env["BROKER_DB_PATH"] = os.path.join(self.tmp, "ttl.db")
        env["BROKER_AUDIT_LOG_PATH"] = os.path.join(self.tmp, "ttl.audit.log")
        env["BROKER_TOKEN_TTL"] = "60"
        with contextlib.redirect_stdout(io.StringIO()):
            broker = coreci_broker.boot(dict(env))
        self.addCleanup(broker.audit.close)
        self.addCleanup(broker.store.close)
        _seed_direct(broker, self.app_uuid, self.user_id)
        server = coreci_broker.serve("127.0.0.1", 0, broker)
        self.addCleanup(server.shutdown)
        self.addCleanup(server.server_close)
        base = "http://127.0.0.1:%d" % server.server_address[1]
        payload = {"agent_id": "ttl-agent", "app_registration_uuid": self.app_uuid,
                   "user_id": self.user_id, "tags": []}
        status, body = self._request("POST", REGISTER_PATH, payload, base=base)
        self.assertEqual(201, status)
        delta = self._ttl_delta_seconds(body)
        self.assertGreater(delta, 30)
        self.assertLess(delta, 90)
        status, body = self._request("POST", REGISTER_PATH, payload, base=base)
        self.assertEqual(200, status)
        delta = self._ttl_delta_seconds(body)
        self.assertGreater(delta, 30)
        self.assertLess(delta, 90)

    def test_user_allowlist_rejects_slash_backslash_and_control_chars(self):
        for bad in ("a/b", "a\\b", "a\nb", "a\x00b", "a\x1fb", "a\x7fb"):
            status, body = self._request("POST", USERS_PATH, {"user_id": bad}, self.admin)
            self.assertEqual(400, status, repr(bad))
            self.assertEqual({"error": 400, "message": "invalid user_id"}, body)
        status, body = self._request("GET", USERS_PATH, headers=self.admin)
        self.assertEqual([self.user_id], [row["user_id"] for row in body["allowlist"]])

    def test_bearer_scheme_with_valid_creds_is_401(self):
        header = "Bearer " + base64.b64encode(
            ("%s:%s" % (ADMIN_USER, ADMIN_PASS)).encode("utf-8")
        ).decode("ascii")
        status, body = self._request("GET", WHOAMI_PATH, headers={"Authorization": header})
        self.assertEqual(401, status)
        self.assertEqual(401, body["error"])

    def test_basic_without_colon_is_401(self):
        header = "Basic " + base64.b64encode(b"just-a-username").decode("ascii")
        status, body = self._request("GET", WHOAMI_PATH, headers={"Authorization": header})
        self.assertEqual(401, status)
        self.assertEqual(401, body["error"])

    def test_basic_garbage_rejected_lowercase_scheme_accepted(self):
        status, body = self._request(
            "GET", WHOAMI_PATH, headers={"Authorization": "Basic !!!not-base64!!!"}
        )
        self.assertEqual(401, status)
        header = _basic_header(ADMIN_USER, ADMIN_PASS).replace("Basic", "basic")
        status, body = self._request("GET", WHOAMI_PATH, headers={"Authorization": header})
        self.assertEqual(200, status)
        self.assertEqual(ADMIN_USER, body["username"])


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