FabledSteward/tests/integration/test_api_tokens.py

"""Integration: personal API tokens against real Postgres.

Resolution reads the token and its user's CURRENT role in one query, so a
demoted or disabled user's tokens weaken or die with nothing revoked by hand.
The routes are driven through the test client so their transactions run for
real: a mocked session cannot catch a double-begin (#884).
"""
from __future__ import annotations

import asyncio
import os
import re
import time
import uuid
from datetime import datetime, timedelta, timezone

import pytest

pytestmark = pytest.mark.integration

_NEEDS_DB = pytest.mark.skipif(
    not os.environ.get("STEWARD_DATABASE_URL"),
    reason="integration test needs a live Postgres (STEWARD_DATABASE_URL)",
)


@pytest.fixture
def app():
    if not os.environ.get("STEWARD_DATABASE_URL"):
        pytest.skip("needs Postgres")
    from steward.app import create_app
    return create_app(testing=False)


async def _make_user(app, role):
    from steward.models.users import User
    uid = str(uuid.uuid4())
    async with app.db_sessionmaker() as s:
        async with s.begin():
            s.add(User(id=uid, username=f"tok-{uid[:8]}", email=f"{uid}@test",
                       password_hash="x", role=role, is_active=True))
    return uid


async def _mint(app, uid, user_role, role, **kw):
    from steward.core import api_tokens
    async with app.db_sessionmaker() as s:
        raw, token = await api_tokens.create_token(
            s, user_id=uid, user_role=user_role, name="it", role=role, **kw)
        await s.commit()
    return raw, token.id


@_NEEDS_DB
def test_resolve_lifecycle(app):
    from sqlalchemy import select, update
    from steward.core import api_tokens
    from steward.models.api_tokens import ApiToken
    from steward.models.users import User, UserRole

    async def _go():
        uid = await _make_user(app, UserRole.admin)
        raw, tid = await _mint(app, uid, UserRole.admin, UserRole.operator)

        caller = await api_tokens.resolve_token(app, raw)
        assert caller is not None
        assert (caller.user_id, caller.role, caller.token_id) == (uid, UserRole.operator, tid)
        async with app.db_sessionmaker() as s:
            used = (await s.execute(select(ApiToken.last_used_at).where(ApiToken.id == tid))).scalar_one()
        assert used is not None

        # A wrong token, and the right one with a character changed, resolve to nobody.
        assert await api_tokens.resolve_token(app, api_tokens.generate_token()) is None
        assert await api_tokens.resolve_token(app, raw[:-1] + ("A" if raw[-1] != "A" else "B")) is None

        # Demoting the user caps the token at the user's new role.
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await s.execute(update(User).where(User.id == uid).values(role=UserRole.viewer))
        assert (await api_tokens.resolve_token(app, raw)).role == UserRole.viewer

        # Disabling the user kills the token.
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await s.execute(update(User).where(User.id == uid).values(is_active=False))
        assert await api_tokens.resolve_token(app, raw) is None
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await s.execute(update(User).where(User.id == uid).values(is_active=True))
        assert await api_tokens.resolve_token(app, raw) is not None

        # Revoking kills it, and only the owner can revoke.
        other = await _make_user(app, UserRole.admin)
        async with app.db_sessionmaker() as s:
            assert await api_tokens.revoke_token(s, user_id=other, token_id=tid) is None
            await s.commit()
        assert await api_tokens.resolve_token(app, raw) is not None
        async with app.db_sessionmaker() as s:
            assert await api_tokens.revoke_token(s, user_id=uid, token_id=tid) is not None
            await s.commit()
        assert await api_tokens.resolve_token(app, raw) is None

        # An expired token resolves to nobody.
        raw2, tid2 = await _mint(app, uid, UserRole.viewer, UserRole.viewer)
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await s.execute(update(ApiToken).where(ApiToken.id == tid2).values(
                    expires_at=datetime.now(timezone.utc) - timedelta(seconds=1)))
        assert await api_tokens.resolve_token(app, raw2) is None

        # Deleting the user takes their tokens with them.
        async with app.db_sessionmaker() as s:
            async with s.begin():
                await s.execute(User.__table__.delete().where(User.id == uid))
            left = (await s.execute(select(ApiToken).where(ApiToken.user_id == uid))).all()
        assert left == []

    asyncio.run(_go())


@_NEEDS_DB
def test_routes_create_list_and_revoke(app):
    from sqlalchemy import select
    from steward.core import api_tokens
    from steward.models.api_tokens import ApiToken
    from steward.models.audit import AuditEvent
    from steward.models.users import UserRole

    async def _go():
        uid = await _make_user(app, UserRole.operator)
        client = app.test_client()
        async with client.session_transaction() as sess:
            sess["user_id"] = uid
            sess["user_role"] = "operator"
            sess["username"] = "route-test"
            sess["_csrf"] = "test-token"
            sess["login_at"] = time.time()
        csrf = {"X-CSRF-Token": "test-token"}

        resp = await client.get("/ai-access/")
        assert resp.status_code == 200
        page = await resp.get_data(as_text=True)
        assert "You have no tokens yet" in page
        # An operator can't offer admin.
        assert 'value="admin"' not in page
        # The connect card names the endpoint, with a placeholder for the token.
        assert 'value="http://localhost/mcp"' in page and "claude mcp add" in page

        resp = await client.post("/ai-access/tokens", form={
            "name": "laptop", "role": "admin", "expires_in_days": ""}, headers=csrf)
        assert resp.status_code == 400
        assert "higher role" in await resp.get_data(as_text=True)

        resp = await client.post("/ai-access/tokens", form={
            "name": "laptop", "role": "operator", "expires_in_days": "30"}, headers=csrf)
        assert resp.status_code == 200
        created = await resp.get_data(as_text=True)
        raw = re.search(r'value="(stw_[^"]+)"', created).group(1)
        # The recipes on the created page carry the new token, ready to paste.
        assert f"Bearer {raw}" in created and "mcp-remote" in created
        caller = await api_tokens.resolve_token(app, raw)
        assert caller is not None and caller.role == UserRole.operator

        async with app.db_sessionmaker() as s:
            token = (await s.execute(select(ApiToken).where(ApiToken.user_id == uid))).scalar_one()
        page = await (await client.get("/ai-access/")).get_data(as_text=True)
        assert "laptop" in page and token.prefix in page
        assert raw not in page

        resp = await client.post(f"/ai-access/tokens/{token.id}/revoke", headers=csrf)
        assert resp.status_code == 302
        assert await api_tokens.resolve_token(app, raw) is None

        async with app.db_sessionmaker() as s:
            actions = (await s.execute(
                select(AuditEvent.action).where(AuditEvent.user_id == uid))).scalars().all()
        assert "api_token.create" in actions and "api_token.revoke" in actions

    asyncio.run(_go())