xerrors--yuxi/backend/package/yuxi/services/auth_service.py

"""认证服务。"""

import hashlib
import secrets
import string
from dataclasses import dataclass
from datetime import timedelta

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from yuxi.storage.postgres.models_business import APIKey, CLIAuthSession, Department, User
from yuxi.utils.auth_utils import AuthUtils
from yuxi.utils.datetime_utils import utc_now_naive

CLI_AUTH_SESSION_TTL_SECONDS = 10 * 60
CLI_AUTH_POLL_INTERVAL_SECONDS = 2
CLI_AUTH_DEFAULT_KEY_NAME = "Yuxi CLI"
CLI_AUTH_USER_CODE_ALPHABET = "".join(ch for ch in string.ascii_uppercase + string.digits if ch not in "0O1I")

CLI_AUTH_STATUS_PENDING = "pending"
CLI_AUTH_STATUS_APPROVED = "approved"
CLI_AUTH_STATUS_CONSUMED = "consumed"
CLI_AUTH_STATUS_EXPIRED = "expired"


@dataclass
class CLIAuthError(Exception):
    code: str
    message: str
    status_code: int = 400


def _hash_secret(value: str) -> str:
    return hashlib.sha256(value.encode()).hexdigest()


def _generate_device_code() -> str:
    return f"yxcli_{secrets.token_urlsafe(32)}"


def _generate_user_code() -> str:
    raw = "".join(secrets.choice(CLI_AUTH_USER_CODE_ALPHABET) for _ in range(8))
    return f"{raw[:4]}-{raw[4:]}"


async def _generate_unique_user_code(db: AsyncSession) -> str:
    for _ in range(10):
        user_code = _generate_user_code()
        result = await db.execute(select(CLIAuthSession.id).filter(CLIAuthSession.user_code == user_code))
        if result.scalar_one_or_none() is None:
            return user_code
    raise RuntimeError("无法生成唯一 CLI 授权码")


def _expire_if_needed(session: CLIAuthSession, now=None) -> bool:
    now = now or utc_now_naive()
    if session.status in {CLI_AUTH_STATUS_PENDING, CLI_AUTH_STATUS_APPROVED} and session.expires_at <= now:
        session.status = CLI_AUTH_STATUS_EXPIRED
        return True
    return False


async def create_cli_auth_session(db: AsyncSession, key_name: str | None = None) -> tuple[CLIAuthSession, str]:
    device_code = _generate_device_code()
    now = utc_now_naive()
    session = CLIAuthSession(
        device_code_hash=_hash_secret(device_code),
        user_code=await _generate_unique_user_code(db),
        status=CLI_AUTH_STATUS_PENDING,
        key_name=(key_name or CLI_AUTH_DEFAULT_KEY_NAME).strip() or CLI_AUTH_DEFAULT_KEY_NAME,
        created_at=now,
        expires_at=now + timedelta(seconds=CLI_AUTH_SESSION_TTL_SECONDS),
    )
    db.add(session)
    await db.commit()
    await db.refresh(session)
    return session, device_code


async def get_cli_auth_session_for_user(
    db: AsyncSession, user_code: str, *, for_update: bool = False
) -> CLIAuthSession:
    stmt = select(CLIAuthSession).filter(CLIAuthSession.user_code == user_code.strip().upper())
    if for_update:
        stmt = stmt.with_for_update()
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()
    if session is None:
        raise CLIAuthError("not_found", "授权会话不存在", status_code=404)
    if _expire_if_needed(session):
        await db.commit()
        raise CLIAuthError("expired_token", "授权会话已过期", status_code=410)
    return session


async def approve_cli_auth_session(db: AsyncSession, user_code: str, user: User) -> CLIAuthSession:
    session = await get_cli_auth_session_for_user(db, user_code, for_update=True)
    if session.status == CLI_AUTH_STATUS_CONSUMED:
        raise CLIAuthError("already_consumed", "授权会话已完成", status_code=409)
    if session.status == CLI_AUTH_STATUS_APPROVED:
        raise CLIAuthError("already_approved", "授权会话已批准", status_code=409)
    if session.status != CLI_AUTH_STATUS_PENDING:
        raise CLIAuthError("invalid_state", "授权会话状态无效", status_code=409)

    session.status = CLI_AUTH_STATUS_APPROVED
    session.approved_user_id = user.id
    session.approved_at = utc_now_naive()
    await db.commit()
    await db.refresh(session)
    return session


async def exchange_cli_auth_token(db: AsyncSession, device_code: str) -> dict:
    result = await db.execute(
        select(CLIAuthSession).filter(CLIAuthSession.device_code_hash == _hash_secret(device_code)).with_for_update()
    )
    session = result.scalar_one_or_none()
    if session is None:
        raise CLIAuthError("invalid_request", "授权会话不存在", status_code=404)
    if _expire_if_needed(session):
        await db.commit()
        raise CLIAuthError("expired_token", "授权会话已过期", status_code=410)
    if session.status == CLI_AUTH_STATUS_PENDING:
        raise CLIAuthError("authorization_pending", "等待浏览器授权", status_code=400)
    if session.status == CLI_AUTH_STATUS_CONSUMED:
        raise CLIAuthError("already_consumed", "授权会话已完成", status_code=409)
    if session.status != CLI_AUTH_STATUS_APPROVED or not session.approved_user_id:
        raise CLIAuthError("invalid_state", "授权会话状态无效", status_code=409)

    user_result = await db.execute(
        select(User, Department.name)
        .outerjoin(Department, User.department_id == Department.id)
        .filter(User.id == session.approved_user_id, User.is_deleted == 0)
    )
    row = user_result.one_or_none()
    if row is None:
        raise CLIAuthError("invalid_user", "授权用户不存在", status_code=409)
    user, department_name = row

    full_key, key_hash, key_prefix = AuthUtils.generate_api_key()
    api_key = APIKey(
        key_hash=key_hash,
        key_prefix=key_prefix,
        name=session.key_name,
        user_id=user.id,
        created_by=str(user.id),
    )
    db.add(api_key)
    await db.flush()

    session.status = CLI_AUTH_STATUS_CONSUMED
    session.api_key_id = api_key.id
    session.consumed_at = utc_now_naive()
    await db.commit()
    await db.refresh(api_key)

    user_data = user.to_dict()
    user_data["department_name"] = department_name

    return {
        "api_key": api_key.to_dict(),
        "secret": full_key,
        "user": user_data,
    }