broker/tests/test_agent_meta.py

import base64
import contextlib
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"
META_KEYS = {"agent_id", "app_registration_uuid", "user_id", "tags",
             "token_expires_at", "first_registered_at", "last_seen_at"}
FORBIDDEN_RAW = ('"token"', '"token_hash"', '"pepper"')


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


class AgentMetaTests(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.registered = self._register("agent-1", ["b", "a"])
        self.token = self.registered["token"]
        self.pepper = self.registered["pepper"]
        self.app = {"Authorization": "App %s %s" % (self.app_uuid, self.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:
                raw = response.read().decode("utf-8")
                return response.status, json.loads(raw), raw
        except urllib.error.HTTPError as error:
            raw = error.read().decode("utf-8")
            return error.code, json.loads(raw), raw

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

    def _meta(self, agent_id, headers=None):
        return self._request(
            "GET", "/api/broker/agents/" + agent_id, headers=headers if headers is not None else self.app
        )

    def test_identity_match_returns_metadata_only(self):
        status, body, raw = self._meta("agent-1")
        self.assertEqual(200, status)
        self.assertEqual(META_KEYS, set(body))
        self.assertEqual("agent-1", body["agent_id"])
        self.assertEqual(self.app_uuid, body["app_registration_uuid"])
        self.assertEqual(self.user_id, body["user_id"])
        self.assertEqual(self.registered["token_expires_at"], body["token_expires_at"])
        row = self.broker.store.get_agent("agent-1", self.app_uuid)
        self.assertEqual(row[6], body["first_registered_at"])
        self.assertEqual(row[7], body["last_seen_at"])
        for forbidden in FORBIDDEN_RAW:
            self.assertNotIn(forbidden, raw)
        self.assertNotIn(self.token, raw)
        self.assertNotIn(self.pepper, raw)

    def test_tags_roundtrip_is_sorted_canonical_list(self):
        status, body, _ = self._meta("agent-1")
        self.assertEqual(200, status)
        self.assertEqual(["a", "b"], body["tags"])

    def test_url_agent_id_mismatch_is_403(self):
        self._register("agent-2", ["coreci"])
        status, body, raw = self._meta("agent-2")
        self.assertEqual(403, status)
        self.assertEqual(403, body["error"])
        self.assertTrue(body["message"])
        for forbidden in FORBIDDEN_RAW:
            self.assertNotIn(forbidden, raw)

    def test_unknown_agent_id_is_404(self):
        status, body, _ = self._meta("ghost-agent")
        self.assertEqual(404, status)
        self.assertEqual({"error": 404, "message": "not found"}, body)

    def test_missing_app_header_is_401(self):
        status, body, _ = self._meta("agent-1", headers={})
        self.assertEqual(401, status)
        self.assertEqual({"error": 401, "message": "invalid app credentials"}, body)


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