Neo-ZQYY/apps/backend/app/auth/dependencies.py

"""
FastAPI 依赖注入：从 JWT 提取当前用户信息。

用法：
    @router.get("/protected")
    async def protected_endpoint(user: CurrentUser = Depends(get_current_user)):
        print(user.user_id, user.site_id)

    # 允许 pending 用户（受限令牌）访问
    @router.get("/apply")
    async def apply_endpoint(user: CurrentUser = Depends(get_current_user_or_limited)):
        if user.limited:
            ...  # 受限逻辑
"""

from dataclasses import dataclass, field

from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError

from app.auth.jwt import decode_access_token

# Bearer token 提取器
_bearer_scheme = HTTPBearer(auto_error=True)


@dataclass(frozen=True)
class CurrentUser:
    """从 JWT 解析出的当前用户上下文。"""

    user_id: int
    site_id: int = 0
    roles: list[str] = field(default_factory=list)
    status: str = "pending"
    limited: bool = False


async def get_current_user(
    credentials: HTTPAuthorizationCredentials = Depends(_bearer_scheme),
) -> CurrentUser:
    """
    FastAPI 依赖：从 Authorization header 提取 JWT，验证后返回用户信息。

    要求完整令牌（非 limited），失败时抛出 401。
    """
    token = credentials.credentials
    try:
        payload = decode_access_token(token)
    except JWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="无效的令牌",
            headers={"WWW-Authenticate": "Bearer"},
        )

    # 受限令牌不允许通过此依赖
    if payload.get("limited"):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="受限令牌无法访问此端点",
            headers={"WWW-Authenticate": "Bearer"},
        )

    user_id_raw = payload.get("sub")
    site_id = payload.get("site_id")

    if user_id_raw is None or site_id is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="令牌缺少必要字段",
            headers={"WWW-Authenticate": "Bearer"},
        )

    try:
        user_id = int(user_id_raw)
    except (TypeError, ValueError):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="令牌中 user_id 格式无效",
            headers={"WWW-Authenticate": "Bearer"},
        )

    roles = payload.get("roles", [])

    return CurrentUser(
        user_id=user_id,
        site_id=site_id,
        roles=roles,
        status="approved",
        limited=False,
    )


async def get_current_user_or_limited(
    credentials: HTTPAuthorizationCredentials = Depends(_bearer_scheme),
) -> CurrentUser:
    """
    FastAPI 依赖：允许 pending 用户（受限令牌）访问。

    - 受限令牌（limited=True）：返回 CurrentUser(limited=True, roles=[], status="pending")
    - 完整令牌：正常返回 CurrentUser
    """
    token = credentials.credentials
    try:
        payload = decode_access_token(token)
    except JWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="无效的令牌",
            headers={"WWW-Authenticate": "Bearer"},
        )

    user_id_raw = payload.get("sub")
    if user_id_raw is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="令牌缺少必要字段",
            headers={"WWW-Authenticate": "Bearer"},
        )

    try:
        user_id = int(user_id_raw)
    except (TypeError, ValueError):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="令牌中 user_id 格式无效",
            headers={"WWW-Authenticate": "Bearer"},
        )

    # 受限令牌：pending 用户
    if payload.get("limited"):
        return CurrentUser(
            user_id=user_id,
            site_id=0,
            roles=[],
            status="pending",
            limited=True,
        )

    # 完整令牌：要求 site_id
    site_id = payload.get("site_id")
    if site_id is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="令牌缺少必要字段",
            headers={"WWW-Authenticate": "Bearer"},
        )

    roles = payload.get("roles", [])

    return CurrentUser(
        user_id=user_id,
        site_id=site_id,
        roles=roles,
        status="approved",
        limited=False,
    )