AgentLab/backend/tests/fakes/auth_env.py

"""身份相关测试的公共部分：把登录设置、用户、组、会话、审计清回关登录的样子，建用户，登录拿一个带会话的客户端。

登录设置是进程内缓存 + settings 表的一行：一个测试开了登录不复原，同一个进程里后面的测试全会 401。fixture 一律在
finally 里调 reset()。
"""
from __future__ import annotations

from typing import Any

from httpx import ASGITransport, AsyncClient
from sqlalchemy import delete

from app.db.base import SessionLocal

#: 开了登录时非 GET 请求必须带的头
CSRF = {"X-AgentLab-Request": "1"}


async def reset() -> None:
    from app.auth import SETTING_KEY, keys, remember_settings, sessions
    from app.db.models import AuditEvent, AuthSession, Group, GroupMember, Setting, User
    from app.db.models_apps import ApiReceipt
    from app.db.models_identity import AccessKey, AccessKeyRequest

    async with SessionLocal() as session:
        for model in (AccessKeyRequest, ApiReceipt, AccessKey, AuthSession, GroupMember, Group, User, AuditEvent):
            await session.execute(delete(model))
        await session.execute(delete(Setting).where(Setting.key == SETTING_KEY))
        await session.commit()
    remember_settings({})
    sessions._failures.clear()
    # 访问密钥的进程内状态：限流的桶、还没写库的请求记录
    keys._BUCKETS.clear()
    keys._PENDING.clear()


def client(**headers: str) -> AsyncClient:
    from app.main import app

    return AsyncClient(transport=ASGITransport(app=app), base_url="http://test", headers={**CSRF, **headers})


async def make_user(username: str, role: str, password: str = "password-123", **extra: Any) -> Any:
    from app.auth.passwords import hash_password
    from app.db.models import User

    async with SessionLocal() as session:
        user = User(username=username, display_name=username, role=role, password_hash=hash_password(password),
                    **extra)
        session.add(user)
        await session.commit()
        await session.refresh(user)
        return user


async def turn_on_local(admin: str = "root", password: str = "password-123") -> AsyncClient:
    """关登录 → 账号密码登录，同一个请求建出第一个管理员；交回已经以管理员登录的客户端。"""
    c = client()
    r = await c.put("/api/auth/settings", json={"mode": "local",
                                                 "bootstrap": {"username": admin, "password": password}})
    assert r.status_code == 200, r.text
    return c


async def login(username: str, password: str = "password-123") -> AsyncClient:
    c = client()
    r = await c.post("/api/auth/login", json={"username": username, "password": password})
    assert r.status_code == 200, r.text
    return c


def keyed(secret: str, **headers: str) -> AsyncClient:
    """带访问密钥的客户端：只有 Authorization 头，不带 CSRF 头、不带 cookie（外部系统就是这样调的）。"""
    from app.main import app

    return AsyncClient(transport=ASGITransport(app=app), base_url="http://test",
                       headers={"Authorization": f"Bearer {secret}", **headers})


async def make_key(owner: Any, *, scopes: tuple[str, ...] = ("apps.read", "apps.run", "tasks.read", "tasks.respond",
                                                             "reports.read"),
                   app_ids: list[str] | None = None, days: float = 90, name: str = "测试密钥") -> tuple[Any, str]:
    """直接在库里签一把密钥（不经接口），交回 (行, 明文)。owner 是 User 行。"""
    from datetime import timedelta

    from app.auth import keys
    from app.db.base import utcnow
    from app.db.models_identity import AccessKey

    secret = keys.new_secret()
    async with SessionLocal() as session:
        row = AccessKey(owner_user_id=owner.id, name=name, key_hash=keys.key_hash(secret), prefix=keys.prefix_of(secret),
                        scopes=list(scopes), app_ids=app_ids, expires_at=utcnow() + timedelta(days=days))
        session.add(row)
        await session.commit()
        await session.refresh(row)
    return row, secret


async def make_service(username: str = "svc-portal", **extra: Any) -> Any:
    from app.db.models import User

    async with SessionLocal() as session:
        user = User(username=username, display_name=username, role="user", kind="service", **extra)
        session.add(user)
        await session.commit()
        await session.refresh(user)
        return user


async def personal_keys(on: bool = True) -> None:
    """个人访问密钥的开关（登录设置里的 personal_keys），直接改缓存和库。"""
    from app.auth import PERSONAL_KEYS, SETTING_KEY, load_settings, remember_settings
    from app.db.models import Setting

    value = {**await load_settings(force=True), PERSONAL_KEYS: on}
    async with SessionLocal() as session:
        row = await session.get(Setting, SETTING_KEY)
        if row is None:
            session.add(Setting(key=SETTING_KEY, value=value))
        else:
            row.value = value
        await session.commit()
    remember_settings(value)