GenieHive/tests/test_control_authorization.py

from pathlib import Path

import pytest
from fastapi.testclient import TestClient

from geniehive_control.keys import hash_api_key
from geniehive_control.main import create_app
from geniehive_control.models import HostRegistration, RegisteredService
from geniehive_control.upstream import UpstreamClient


class _FakeResponse:
    def __init__(self, payload: dict, status_code: int = 200) -> None:
        self._payload = payload
        self.status_code = status_code
        self.text = str(payload)

    def json(self) -> dict:
        return self._payload


class _FakePoster:
    async def post(self, url: str, *, json: dict, headers: dict[str, str] | None = None) -> _FakeResponse:
        if url.endswith("/v1/embeddings"):
            return _FakeResponse({"object": "list", "data": [{"embedding": [0.1, 0.2]}]})
        return _FakeResponse({"object": "chat.completion", "model": json["model"], "choices": []})


def _write_config(tmp_path: Path, *, static_key: bool = False) -> Path:
    config_path = tmp_path / "control.yaml"
    static_auth = """
  client_api_keys:
    - static-key
""" if static_key else ""
    config_path.write_text(
        f"""
auth:
{static_auth}  enable_named_client_keys: true
authorization:
  enforce_model_allowlists: true
  enforce_operation_allowlists: true
  empty_allowlist_means_no_access: true
storage:
  sqlite_path: "{tmp_path / 'geniehive.sqlite3'}"
"""
    )
    return config_path


def _register_services(app) -> None:
    app.state.registry.register_host(
        HostRegistration(
            host_id="atlas-01",
            address="127.0.0.1",
            services=[
                RegisteredService(
                    service_id="atlas-01/chat/qwen",
                    host_id="atlas-01",
                    kind="chat",
                    endpoint="http://127.0.0.1:18091",
                    assets=[{"asset_id": "archive_migrator", "loaded": True}],
                    state={"health": "healthy", "accept_requests": True},
                    observed={"p50_latency_ms": 100},
                ),
                RegisteredService(
                    service_id="atlas-01/embeddings/bge",
                    host_id="atlas-01",
                    kind="embeddings",
                    endpoint="http://127.0.0.1:18092",
                    assets=[{"asset_id": "bge-small", "loaded": True}],
                    state={"health": "healthy", "accept_requests": True},
                    observed={"p50_latency_ms": 100},
                ),
            ],
        )
    )


def _create_named_key(
    app,
    raw_key: str,
    *,
    allowed_models: list[str],
    allowed_operations: list[str],
) -> None:
    app.state.registry.create_client_key(
        key_id=f"ck_{raw_key}",
        key_hash=hash_api_key(raw_key, secret="test-secret"),
        display_name="Scoped User",
        principal_type="person",
        principal_ref="scoped-user",
        role="developer",
        allowed_models=allowed_models,
        allowed_operations=allowed_operations,
    )


def test_named_key_allows_scoped_chat_request(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("GENIEHIVE_KEY_HASH_SECRET", "test-secret")
    app = create_app(_write_config(tmp_path), upstream_client=UpstreamClient(client=_FakePoster()))
    _register_services(app)
    _create_named_key(
        app,
        "gh_allowed",
        allowed_models=["archive_migrator"],
        allowed_operations=["chat"],
    )
    client = TestClient(app)

    response = client.post(
        "/v1/chat/completions",
        headers={"X-Api-Key": "gh_allowed"},
        json={"model": "archive_migrator", "messages": [{"role": "user", "content": "hello"}]},
    )

    assert response.status_code == 200


def test_named_key_denies_unlisted_operation(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("GENIEHIVE_KEY_HASH_SECRET", "test-secret")
    app = create_app(_write_config(tmp_path), upstream_client=UpstreamClient(client=_FakePoster()))
    _register_services(app)
    _create_named_key(
        app,
        "gh_chat_only",
        allowed_models=["*"],
        allowed_operations=["chat"],
    )
    client = TestClient(app)

    response = client.post(
        "/v1/embeddings",
        headers={"X-Api-Key": "gh_chat_only"},
        json={"model": "bge-small", "input": "hello"},
    )

    assert response.status_code == 403
    assert response.json()["error"]["code"] == "authorization_error"


def test_named_key_denies_unlisted_model(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("GENIEHIVE_KEY_HASH_SECRET", "test-secret")
    app = create_app(_write_config(tmp_path), upstream_client=UpstreamClient(client=_FakePoster()))
    _register_services(app)
    _create_named_key(
        app,
        "gh_archive_only",
        allowed_models=["archive_migrator"],
        allowed_operations=["chat"],
    )
    client = TestClient(app)

    response = client.post(
        "/v1/chat/completions",
        headers={"X-Api-Key": "gh_archive_only"},
        json={"model": "other_role", "messages": [{"role": "user", "content": "hello"}]},
    )

    assert response.status_code == 403
    assert response.json()["error"]["code"] == "authorization_error"


def test_empty_allowlist_denies_when_configured(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("GENIEHIVE_KEY_HASH_SECRET", "test-secret")
    app = create_app(_write_config(tmp_path), upstream_client=UpstreamClient(client=_FakePoster()))
    _register_services(app)
    _create_named_key(
        app,
        "gh_empty",
        allowed_models=[],
        allowed_operations=[],
    )
    client = TestClient(app)

    response = client.post(
        "/v1/chat/completions",
        headers={"X-Api-Key": "gh_empty"},
        json={"model": "archive_migrator", "messages": [{"role": "user", "content": "hello"}]},
    )

    assert response.status_code == 403


def test_static_key_is_not_restricted_by_named_key_allowlists(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("GENIEHIVE_KEY_HASH_SECRET", "test-secret")
    app = create_app(
        _write_config(tmp_path, static_key=True),
        upstream_client=UpstreamClient(client=_FakePoster()),
    )
    _register_services(app)
    client = TestClient(app)

    response = client.post(
        "/v1/embeddings",
        headers={"X-Api-Key": "static-key"},
        json={"model": "bge-small", "input": "hello"},
    )

    assert response.status_code == 200