broker/tests/test_app_auth.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"
PROTECTED_PATH = "/test/app-auth"
AUTH_401_BODY = {"error": 401, "message": "invalid app credentials"}


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


def _protected_echo(request):
    request.send_json(200, dict(request.app_auth))


class AppAuthTests(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)
        self.broker.router.add(
            "GET", PROTECTED_PATH, coreci_broker._require_app(_protected_echo)
        )
        self.token = self._register("agent-1")["token"]

    def _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:
                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):
        payload = {
            "agent_id": agent_id,
            "app_registration_uuid": self.app_uuid,
            "user_id": self.user_id,
            "tags": ["coreci"],
        }
        status, body = self._request("POST", REGISTER_PATH, payload)
        self.assertEqual(201, status)
        return body

    def _auth(self, authorization):
        headers = {} if authorization is None else {"Authorization": authorization}
        return self._request("GET", PROTECTED_PATH, headers=headers)

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

    def _revoke_app(self):
        status, _ = self._request(
            "DELETE", APPS_PATH + "/" + self.app_uuid, headers=self.admin
        )
        self.assertEqual(200, status)

    def _patch_column(self, agent_id, column, value):
        self.broker.store.conn.execute(
            "UPDATE agent_registrations SET %s = ? WHERE agent_id = ?" % column,
            (value, agent_id),
        )
        self.broker.store.conn.commit()

    def _count_verify_calls(self):
        counter = {"calls": 0}
        original = coreci_broker.verify_token

        def counted(token, stored):
            counter["calls"] += 1
            return original(token, stored)

        coreci_broker.verify_token = counted
        self.addCleanup(setattr, coreci_broker, "verify_token", original)
        return counter

    def _past(self):
        return (datetime.datetime.now(datetime.timezone.utc)
                - datetime.timedelta(seconds=60)).isoformat()

    def _future(self):
        return (datetime.datetime.now(datetime.timezone.utc)
                + datetime.timedelta(seconds=3600)).isoformat()

    def test_valid_header_returns_context_and_updates_last_seen(self):
        past = self._past()
        self._patch_column("agent-1", "last_seen_at", past)
        status, body = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual("agent-1", body["agent_id"])
        self.assertEqual(self.user_id, body["user_id"])
        self.assertEqual(self.app_uuid, body["app_registration_uuid"])
        row = self.broker.store.get_agent("agent-1", self.app_uuid)
        self.assertEqual(row[5], body["token_expires_at"])
        self.assertGreater(row[7], past)

    def test_401_matrix_uniform_contract(self):
        cases = (
            (None),
            (_basic_header(ADMIN_USER, ADMIN_PASS)),
            ("App"),
            ("App %s" % self.app_uuid),
            ("App %s %s extra" % (self.app_uuid, self.token)),
            (self._app_header(app_uuid=str(uuidlib.uuid4()))),
            (self._app_header(token="A" * 44)),
            ("App not-a-uuid %s" % self.token),
            ("App %s %s" % (self.app_uuid.upper(), self.token)),
        )
        for header in cases:
            with self.subTest(header=header):
                status, body = self._auth(header)
                self.assertEqual(401, status)
                self.assertEqual(AUTH_401_BODY, body)

    def test_cache_first_success_scans_then_identical_call_hits_cache(self):
        counter = self._count_verify_calls()
        status, body = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual(1, counter["calls"])
        status, body = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual("agent-1", body["agent_id"])
        self.assertEqual(1, counter["calls"])
        self.assertIn((self.app_uuid, self.token), self.broker.token_cache)

    def test_clear_token_cache_forces_reverify(self):
        counter = self._count_verify_calls()
        self._auth(self._app_header())
        self.assertEqual(1, counter["calls"])
        self.broker.clear_token_cache()
        self.assertEqual({}, self.broker.token_cache)
        status, _ = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual(2, counter["calls"])

    def test_cache_hit_revoked_app_still_200_until_ttl(self):
        counter = self._count_verify_calls()
        self._auth(self._app_header())
        self._revoke_app()
        status, body = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual("agent-1", body["agent_id"])
        self.assertEqual(1, counter["calls"])
        row = self.broker.store.get_row("app_registrations", "uuid", self.app_uuid)
        self.assertEqual("revoked", row[2])

    def test_revoked_app_first_auth_still_200(self):
        self._revoke_app()
        counter = self._count_verify_calls()
        status, body = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual("agent-1", body["agent_id"])
        self.assertEqual(1, counter["calls"])

    def test_revoked_app_expired_row_is_401(self):
        self._revoke_app()
        self._patch_column("agent-1", "token_expires_at", self._past())
        status, body = self._auth(self._app_header())
        self.assertEqual(401, status)
        self.assertEqual(AUTH_401_BODY, body)

    def test_cache_hit_expired_row_401_then_evicted_and_reverifies(self):
        counter = self._count_verify_calls()
        self._auth(self._app_header())
        self.assertEqual(1, counter["calls"])
        self._patch_column("agent-1", "token_expires_at", self._past())
        status, body = self._auth(self._app_header())
        self.assertEqual(401, status)
        self.assertEqual(AUTH_401_BODY, body)
        self.assertEqual(1, counter["calls"])
        self.assertEqual({}, self.broker.token_cache)
        self._patch_column("agent-1", "token_expires_at", self._future())
        status, body = self._auth(self._app_header())
        self.assertEqual(200, status)
        self.assertEqual(2, counter["calls"])
        self.assertIn((self.app_uuid, self.token), self.broker.token_cache)

    def test_cache_hit_row_absent_is_401(self):
        self._auth(self._app_header())
        self.broker.store.conn.execute(
            "DELETE FROM agent_registrations WHERE agent_id = ?", ("agent-1",)
        )
        self.broker.store.conn.commit()
        status, body = self._auth(self._app_header())
        self.assertEqual(401, status)
        self.assertEqual(AUTH_401_BODY, body)
        self.assertEqual({}, self.broker.token_cache)

    def test_scrypt_hash_format_authenticates(self):
        token2 = self._register("agent-2")["token"]
        original_bcrypt = coreci_broker.bcrypt
        coreci_broker.bcrypt = None
        try:
            scrypt_hash = coreci_broker.hash_token(token2)
        finally:
            coreci_broker.bcrypt = original_bcrypt
        self.assertTrue(scrypt_hash.startswith("scrypt$"))
        self._patch_column("agent-2", "token_hash", scrypt_hash)
        status, body = self._auth(self._app_header(token=token2))
        self.assertEqual(200, status)
        self.assertEqual("agent-2", body["agent_id"])

    def test_bcrypt_hash_format_authenticates(self):
        if coreci_broker.bcrypt is None:
            self.skipTest("bcrypt not importable in this environment")
        token3 = self._register("agent-3")["token"]
        stored = self.broker.store.get_agent("agent-3", self.app_uuid)[4]
        self.assertTrue(stored.startswith("$2"))
        status, body = self._auth(self._app_header(token=token3))
        self.assertEqual(200, status)
        self.assertEqual("agent-3", body["agent_id"])


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