broker/tests/test_acceptance_p3.py

import base64
import contextlib
import datetime
import hashlib
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

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"
AUDIT_PATH = "/api/broker/audit"
AGENT_PATH = "/api/broker/agents/%s"
STATE_PATH = "/api/broker/agents/%s/state"
AUTH_401_BODY = {"error": 401, "message": "invalid app credentials"}
REGISTER_403_BODY = {"error": 403, "message": "registration not allowed"}


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


def _derive(pepper, app_uuid, user_id, agent_id, tags):
    material = "%s:%s:%s:%s:%s" % (
        pepper, app_uuid, user_id, agent_id, ",".join(sorted(tags))
    )
    digest = hashlib.sha256(material.encode("utf-8")).digest()
    return base64.b64encode(digest).decode("ascii")


class AcceptanceP3Tests(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,
        }
        self.broker, self.base = self._boot()
        self.admin = {"Authorization": _basic_header(ADMIN_USER, ADMIN_PASS)}
        self.app_uuid = str(uuidlib.uuid4())
        self.user_id = "ceo"
        self._seed(self.base, self.app_uuid, self.user_id)

    def _boot(self, extra=None):
        env = dict(self.env)
        env.update(extra or {})
        with contextlib.redirect_stdout(io.StringIO()):
            broker = coreci_broker.boot(env)
        self.addCleanup(broker.audit.close)
        self.addCleanup(broker.store.close)
        server = coreci_broker.serve("127.0.0.1", 0, broker)
        self.addCleanup(server.shutdown)
        self.addCleanup(server.server_close)
        return broker, "http://127.0.0.1:%d" % server.server_address[1]

    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:
                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, base, app_uuid, user_id):
        status, body = self._request(
            "POST", APPS_PATH, {"uuid": app_uuid}, headers=self.admin, base=base
        )
        self.assertEqual(201, status)
        status, body = self._request(
            "POST", USERS_PATH, {"user_id": user_id}, headers=self.admin, base=base
        )
        self.assertEqual(201, status)

    def _register(self, agent_id, tags, base=None, app_uuid=None, user_id=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, base=base)

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

    def _audit_entries(self, base, query):
        status, body = self._request(
            "GET", AUDIT_PATH + query, headers=self.admin, base=base
        )
        self.assertEqual(200, status)
        return body["entries"]

    def test_ac4_derivation_identical_both_sides_stored_as_slow_hash(self):
        agent_id = "px-1"
        tags = ["zeta", "alpha"]
        status, body = self._register(agent_id, tags)
        self.assertEqual(201, status)
        self.assertTrue(body["created"])
        token = body["token"]
        pepper = self.broker.store.get_config("BROKER_INSTANCE_PEPPER")
        self.assertEqual(pepper, body["pepper"])
        broker_side = _derive(pepper, self.app_uuid, self.user_id, agent_id, tags)
        tool_side = _derive(body["pepper"], self.app_uuid, self.user_id, agent_id, tags)
        self.assertEqual(broker_side, token)
        self.assertEqual(tool_side, token)
        self.assertEqual(44, len(token))
        self.assertTrue(token.isascii())
        row = self.broker.store.get_agent(agent_id, self.app_uuid)
        stored = row[4]
        self.assertNotEqual(token, stored)
        if coreci_broker.bcrypt is None:
            self.assertTrue(stored.startswith("scrypt$"))
        else:
            self.assertTrue(stored.startswith(coreci_broker.BCRYPT_PREFIXES))
        self.assertTrue(coreci_broker.verify_token(token, stored))
        self.assertFalse(coreci_broker.verify_token(token, token))
        self.assertEqual(self.user_id, row[2])
        self.assertEqual(json.dumps(sorted(tags)), row[3])
        self.assertEqual(body["token_expires_at"], row[5])

    def test_ac7_revoked_or_unknown_uuid_denied_existing_token_honored(self):
        status, first = self._register("tool-1", ["coreci"])
        self.assertEqual(201, status)
        token = first["token"]
        status, body = self._request(
            "DELETE", APPS_PATH + "/" + self.app_uuid, headers=self.admin
        )
        self.assertEqual(200, status)
        self.assertEqual({"uuid": self.app_uuid, "status": "revoked"}, body)
        status, body = self._register("tool-2", ["coreci"])
        self.assertEqual(403, status)
        self.assertEqual(REGISTER_403_BODY, body)
        self.assertIsNone(self.broker.store.get_agent("tool-2", self.app_uuid))
        status, body = self._register("tool-3", ["coreci"], app_uuid=str(uuidlib.uuid4()))
        self.assertEqual(403, status)
        self.assertEqual(REGISTER_403_BODY, body)
        status, body = self._request(
            "GET", AGENT_PATH % "tool-1", headers=self._app_header(token)
        )
        self.assertEqual(200, status)
        self.assertEqual("tool-1", body["agent_id"])
        self.assertEqual(self.user_id, body["user_id"])
        self.assertEqual(self.app_uuid, body["app_registration_uuid"])
        self.assertEqual(["coreci"], body["tags"])

    def test_ac8_unknown_or_revoked_user_denied(self):
        status, body = self._register("tool-1", ["coreci"], user_id="intern")
        self.assertEqual(403, status)
        self.assertEqual(REGISTER_403_BODY, body)
        status, body = self._request(
            "POST", USERS_PATH, {"user_id": "intern"}, headers=self.admin
        )
        self.assertEqual(201, status)
        status, body = self._register("tool-1", ["coreci"], user_id="intern")
        self.assertEqual(201, status)
        self.assertTrue(body["created"])
        status, body = self._request(
            "DELETE", USERS_PATH + "/intern", headers=self.admin
        )
        self.assertEqual(200, status)
        self.assertEqual({"user_id": "intern", "status": "revoked"}, body)
        status, body = self._register("tool-1", ["coreci"], user_id="intern")
        self.assertEqual(403, status)
        self.assertEqual(REGISTER_403_BODY, body)

    def test_ac10_expiry_401_reregister_refreshes_expiry_and_audits(self):
        broker, base = self._boot({
            "BROKER_DB_PATH": os.path.join(self.tmp, "ttl.db"),
            "BROKER_AUDIT_LOG_PATH": os.path.join(self.tmp, "ttl.audit.log"),
            "BROKER_TOKEN_TTL": "1",
        })
        self._seed(base, self.app_uuid, self.user_id)
        agent_id = "ttl-tool"
        status, first = self._register(agent_id, ["coreci"], base=base)
        self.assertEqual(201, status)
        delta = (datetime.datetime.fromisoformat(first["token_expires_at"])
                 - datetime.datetime.now(datetime.timezone.utc)).total_seconds()
        self.assertGreater(delta, 0)
        self.assertLess(delta, 3)
        status, body = self._request(
            "GET", AGENT_PATH % agent_id,
            headers=self._app_header(first["token"]), base=base,
        )
        self.assertEqual(200, status)
        time.sleep(1.2)
        status, body = self._request(
            "GET", AGENT_PATH % agent_id,
            headers=self._app_header(first["token"]), base=base,
        )
        self.assertEqual(401, status)
        self.assertEqual(AUTH_401_BODY, body)
        status, body = self._register(agent_id, ["coreci"], base=base, user_id="ghost-user")
        self.assertEqual(403, status)
        status, second = self._register(agent_id, ["coreci"], base=base)
        self.assertEqual(200, status)
        self.assertFalse(second["created"])
        self.assertEqual(first["token"], second["token"])
        self.assertIn("pepper", second)
        self.assertGreater(
            datetime.datetime.fromisoformat(second["token_expires_at"]),
            datetime.datetime.fromisoformat(first["token_expires_at"]),
        )
        status, body = self._request(
            "GET", AGENT_PATH % agent_id,
            headers=self._app_header(second["token"]), base=base,
        )
        self.assertEqual(200, status)
        row = broker.store.get_agent(agent_id, self.app_uuid)
        self.assertEqual(second["token_expires_at"], row[5])
        self.assertTrue(coreci_broker.verify_token(second["token"], row[4]))
        entries = self._audit_entries(base, "?action=register")
        self.assertEqual(["ok", "denied", "ok"], [entry["result"] for entry in entries])
        for entry in entries:
            self.assertEqual(agent_id, entry["agent_id"])
            self.assertEqual(self.app_uuid, entry["app_registration_uuid"])
            self.assertEqual(REGISTER_PATH, entry["path"])
        self.assertEqual(self.user_id, entries[0]["user_id"])
        self.assertEqual("ghost-user", entries[1]["user_id"])
        self.assertEqual(self.user_id, entries[2]["user_id"])

    def test_ac12_state_roundtrip_with_identity_enforcement(self):
        status, a = self._register("tool-a", ["coreci", "alpha"])
        self.assertEqual(201, status)
        status, b = self._register("tool-b", ["coreci"])
        self.assertEqual(201, status)
        state = {
            "agent_id": "tool-a",
            "app_registration_uuid": self.app_uuid,
            "user_id": self.user_id,
            "tags": ["alpha", "coreci"],
            "pepper": a["pepper"],
            "current_token": a["token"],
            "token_expires_at": a["token_expires_at"],
            "first_registered_at": "t0",
            "last_seen_at": "t1",
            "payload": {"counter": 7, "notes": ["x", "y"], "nested": {"k": None}},
        }
        status, body = self._request(
            "PUT", STATE_PATH % "tool-a", state, headers=self._app_header(a["token"])
        )
        self.assertEqual(200, status)
        self.assertEqual({"ok": True}, body)
        status, body = self._request(
            "GET", STATE_PATH % "tool-a", headers=self._app_header(a["token"])
        )
        self.assertEqual(200, status)
        self.assertEqual(state, body)
        self.assertEqual(list(state), list(body))
        status, body = self._request(
            "GET", STATE_PATH % "tool-a", headers=self._app_header(b["token"])
        )
        self.assertEqual(403, status)
        self.assertEqual({"error": 403, "message": "agent identity mismatch"}, body)
        status, body = self._request(
            "PUT", STATE_PATH % "tool-a", {"intruder": True},
            headers=self._app_header(b["token"]),
        )
        self.assertEqual(403, status)
        self.assertEqual({"error": 403, "message": "agent identity mismatch"}, body)
        status, body = self._request(
            "PUT", STATE_PATH % "tool-b", {"mine": True},
            headers=self._app_header(b["token"]),
        )
        self.assertEqual(200, status)
        status, body = self._request(
            "GET", STATE_PATH % "tool-b", headers=self._app_header(b["token"])
        )
        self.assertEqual(200, status)
        self.assertEqual({"mine": True}, body)
        status, body = self._request(
            "GET", STATE_PATH % "ghost-agent", headers=self._app_header(a["token"])
        )
        self.assertEqual(404, status)
        self.assertEqual({"error": 404, "message": "not found"}, body)
        status, body = self._request(
            "GET", STATE_PATH % "tool-a", headers=self._app_header(a["token"])
        )
        self.assertEqual(200, status)
        self.assertEqual(state, body)


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