EvoScientist/tests/test_delegation_auth.py

"""Delegation JWT and BFF service-token authentication tests (design doc 7.3).

Covers the正反例 the brief requires: missing/invalid service token,
malformed/expired/wrongly-signed delegation JWTs, missing claims, lifetime
over 60 seconds, wrong issuer/audience, unregistered deployments, scope
escalation, jti replay, and the thread-claim rules for config vs
thread-level routes. ``GET /api/model-registry`` serves as the probe route.
"""

from __future__ import annotations

import hashlib
import time
import uuid

import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from starlette.applications import Starlette
from starlette.testclient import TestClient

from EvoScientist.model_registry.auth import BffAuthenticator
from EvoScientist.model_registry.endpoint_policy import EndpointPolicy
from EvoScientist.model_registry.http_api import (
    ApiServices,
    model_registry_routes,
)
from EvoScientist.model_registry.platform import DelegationPublicKey
from EvoScientist.model_registry.resolver import ModelRegistryResolver
from EvoScientist.model_registry.snapshots import SnapshotService
from EvoScientist.model_registry.store import ModelRuntimeStore

SERVICE_TOKEN = "bff-service-token"

_PRIVATE_KEY = ec.generate_private_key(ec.SECP256R1())
_PRIVATE_PEM = _PRIVATE_KEY.private_bytes(
    serialization.Encoding.PEM,
    serialization.PrivateFormat.PKCS8,
    serialization.NoEncryption(),
)
_PUBLIC_PEM = _PRIVATE_KEY.public_key().public_bytes(
    serialization.Encoding.PEM,
    serialization.PublicFormat.SubjectPublicKeyInfo,
)

_OTHER_KEY = ec.generate_private_key(ec.SECP256R1())
_OTHER_PRIVATE_PEM = _OTHER_KEY.private_bytes(
    serialization.Encoding.PEM,
    serialization.PrivateFormat.PKCS8,
    serialization.NoEncryption(),
)


def _services(store: ModelRuntimeStore, **auth_overrides) -> ApiServices:
    resolver = ModelRegistryResolver(store)
    auth_kwargs = {
        "service_token": SERVICE_TOKEN,
        "service_token_hash": None,
        "delegation_keys": (
            DelegationPublicKey(deployment_id="webui-1", public_key=_PUBLIC_PEM),
        ),
        "jti_store": store,
    }
    auth_kwargs.update(auth_overrides)
    return ApiServices(
        store=store,
        resolver=resolver,
        snapshot_service=SnapshotService(store, resolver),
        endpoint_policy=EndpointPolicy(),
        authenticator=BffAuthenticator(**auth_kwargs),
    )


@pytest.fixture
def store(tmp_path):
    return ModelRuntimeStore(config_dir=tmp_path)


@pytest.fixture
def client(store):
    services = _services(store)
    app = Starlette(routes=model_registry_routes(lambda: services))
    return TestClient(app)


def _claims(**overrides):
    now = int(time.time())
    claims = {
        "iss": "WebUI",
        "aud": "EvoScientist",
        "sub": "admin-1",
        "scopes": ["model_config:read"],
        "deployment_id": "webui-1",
        "iat": now,
        "exp": now + 30,
        "jti": uuid.uuid4().hex,
    }
    claims.update(overrides)
    return claims


def _encode(claims, pem=_PRIVATE_PEM):
    return jwt.encode(claims, pem, algorithm="ES256")


def _headers(claims=None, *, token=SERVICE_TOKEN, pem=_PRIVATE_PEM):
    headers = {}
    if token is not None:
        headers["Authorization"] = f"Bearer {token}"
    if claims is not None:
        headers["X-Evo-Actor"] = _encode(claims, pem)
    return headers


def test_valid_token_and_delegation_pass(client):
    response = client.get("/api/model-registry", headers=_headers(_claims()))
    assert response.status_code == 200


def test_config_route_does_not_require_thread_claim(client):
    claims = _claims()
    assert "thread_id" not in claims
    response = client.get("/api/model-registry", headers=_headers(claims))
    assert response.status_code == 200


def test_missing_authorization_header_is_401(client):
    response = client.get("/api/model-registry", headers={"X-Evo-Actor": "x"})
    assert response.status_code == 401
    body = response.json()
    assert body["code"] == "UNAUTHENTICATED"
    assert set(body) == {"code", "message", "details", "request_id"}
    assert body["request_id"]


def test_wrong_service_token_is_401(client):
    response = client.get(
        "/api/model-registry",
        headers=_headers(_claims(), token="not-the-token"),
    )
    assert response.status_code == 401
    assert response.json()["code"] == "UNAUTHENTICATED"


def test_service_token_hash_configuration_passes(store):
    services = _services(
        store,
        service_token=None,
        service_token_hash=hashlib.sha256(SERVICE_TOKEN.encode()).hexdigest(),
    )
    app = Starlette(routes=model_registry_routes(lambda: services))
    hashed_client = TestClient(app)
    assert (
        hashed_client.get(
            "/api/model-registry", headers=_headers(_claims())
        ).status_code
        == 200
    )
    assert (
        hashed_client.get(
            "/api/model-registry", headers=_headers(_claims(), token="wrong")
        ).status_code
        == 401
    )


def test_missing_delegation_jwt_is_401(client):
    response = client.get(
        "/api/model-registry",
        headers={"Authorization": f"Bearer {SERVICE_TOKEN}"},
    )
    assert response.status_code == 401


def test_malformed_delegation_jwt_is_401(client):
    response = client.get(
        "/api/model-registry",
        headers=_headers(None) | {"X-Evo-Actor": "not-a-jwt"},
    )
    assert response.status_code == 401


@pytest.mark.parametrize(
    "missing", ["sub", "scopes", "deployment_id", "iat", "exp", "jti"]
)
def test_missing_required_claim_is_401(client, missing):
    claims = _claims()
    del claims[missing]
    response = client.get("/api/model-registry", headers=_headers(claims))
    assert response.status_code == 401


def test_expired_delegation_jwt_is_401(client):
    now = int(time.time())
    claims = _claims(iat=now - 90, exp=now - 60)
    response = client.get("/api/model-registry", headers=_headers(claims))
    assert response.status_code == 401


def test_lifetime_over_60_seconds_is_401(client):
    now = int(time.time())
    claims = _claims(iat=now, exp=now + 120)
    response = client.get("/api/model-registry", headers=_headers(claims))
    assert response.status_code == 401


def test_wrong_signature_is_401(client):
    response = client.get(
        "/api/model-registry",
        headers=_headers(_claims(), pem=_OTHER_PRIVATE_PEM),
    )
    assert response.status_code == 401


def test_wrong_issuer_is_401(client):
    response = client.get(
        "/api/model-registry", headers=_headers(_claims(iss="not-webui"))
    )
    assert response.status_code == 401


def test_wrong_audience_is_401(client):
    response = client.get(
        "/api/model-registry", headers=_headers(_claims(aud="someone-else"))
    )
    assert response.status_code == 401


def test_unregistered_deployment_is_401(client):
    response = client.get(
        "/api/model-registry", headers=_headers(_claims(deployment_id="webui-x"))
    )
    assert response.status_code == 401


def test_insufficient_scope_is_403(client):
    claims = _claims(scopes=["model:select"])
    response = client.get("/api/model-registry", headers=_headers(claims))
    assert response.status_code == 403
    body = response.json()
    assert body["code"] == "FORBIDDEN"
    assert set(body) == {"code", "message", "details", "request_id"}


def test_jti_replay_is_401_delegation_replayed(client):
    headers = _headers(_claims())
    first = client.get("/api/model-registry", headers=headers)
    assert first.status_code == 200
    second = client.get("/api/model-registry", headers=headers)
    assert second.status_code == 401
    assert second.json()["code"] == "DELEGATION_REPLAYED"


def test_invalid_scopes_claim_is_401(client):
    response = client.get(
        "/api/model-registry", headers=_headers(_claims(scopes="model_config:read"))
    )
    assert response.status_code == 401