LabelingSystem/backend/middleware/auth_middleware.py

"""
Authentication Middleware for SSO token verification.
Validates SSO tokens via the SSO center's userinfo endpoint,
with an in-memory cache to reduce external calls.

Also supports admin tokens generated by generate_admin_token.py script.
"""
import logging
from datetime import datetime, timezone
from fastapi import Request, HTTPException, status
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from services.token_cache_service import TokenCacheService
from services.oauth_service import OAuthService
from config import settings
from database import get_db_connection

logger = logging.getLogger(__name__)

# 全局 token 缓存实例
# SSO token 有效期 600 秒，缓存设置为 550 秒（留 50 秒余量）
token_cache = TokenCacheService(
    ttl_seconds=getattr(settings, 'TOKEN_CACHE_TTL', 550)
)


def verify_admin_token(token: str) -> dict:
    """
    验证管理员 Token（从数据库查询）

    Args:
        token: Token 字符串

    Returns:
        dict: 用户信息字典，或 None（Token 无效或已过期）
    """
    try:
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT at.user_id, u.username, u.email, u.role, at.expires_at
                FROM admin_tokens at
                JOIN users u ON at.user_id = u.id
                WHERE at.token = %s AND at.expires_at > %s
            """, (token, datetime.now(timezone.utc)))
            row = cursor.fetchone()

            if not row:
                return None

            return {
                "id": row["user_id"],
                "username": row["username"],
                "email": row["email"],
                "role": row["role"],
                "is_admin_token": True,
            }
    except Exception as e:
        logger.error(f"验证管理员 Token 失败：{e}")
        return None


class AuthMiddleware(BaseHTTPMiddleware):
    """
    SSO Token 认证中间件。
    先查本地缓存，未命中则调用 SSO userinfo 端点验证。
    """

    PUBLIC_PATHS = {
        "/",
        "/health",
        "/docs",
        "/openapi.json",
        "/redoc",
        "/api/oauth/status",
        "/api/oauth/login",
        "/api/oauth/callback",
        "/api/oauth/refresh",
    }

    async def dispatch(self, request: Request, call_next):
        # Skip authentication for public paths
        logger.debug(f"AuthMiddleware: path={request.url.path}, method={request.method}")
        if request.url.path in self.PUBLIC_PATHS:
            logger.debug(f"Skipping auth for public path: {request.url.path}")
            return await call_next(request)

        # Skip authentication for OPTIONS requests (CORS preflight)
        if request.method == "OPTIONS":
            return await call_next(request)

        # Check if OAuth/SSO is enabled
        if not settings.OAUTH_ENABLED:
            return JSONResponse(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                content={
                    "detail": "SSO 认证未配置",
                    "error_type": "sso_not_configured"
                }
            )

        # Extract token from Authorization header
        auth_header = request.headers.get("Authorization")

        if not auth_header:
            return JSONResponse(
                status_code=status.HTTP_401_UNAUTHORIZED,
                content={
                    "detail": "缺少认证令牌",
                    "error_type": "missing_token"
                }
            )

        # Verify Bearer token format
        parts = auth_header.split()
        if len(parts) != 2 or parts[0].lower() != "bearer":
            return JSONResponse(
                status_code=status.HTTP_401_UNAUTHORIZED,
                content={
                    "detail": "无效的认证令牌格式",
                    "error_type": "invalid_token_format"
                }
            )

        sso_token = parts[1]

        try:
            # 1. 先检查是否是管理员 Token（以 admin_token_ 开头）
            user_info = None
            if sso_token.startswith("admin_token_"):
                logger.debug("检测到管理员 Token，尝试从数据库验证")
                user_info = verify_admin_token(sso_token)
                if user_info:
                    logger.info(f"管理员 Token 验证成功：{user_info['username']}")

            # 2. 如果不是管理员 Token，查本地缓存（SSO token）
            if user_info is None:
                user_info = token_cache.get(sso_token)

            # 3. 缓存未命中，调 SSO profile 验证（含角色信息）
            if user_info is None:
                user_info = await OAuthService.verify_sso_token(sso_token)

                # 3. 同步用户到本地数据库（更新角色），获取本地用户ID
                try:
                    local_user = OAuthService.sync_user_from_oauth(oauth_user_info=user_info)
                    # 将本地user.id也存入user_info，供后续使用
                    user_info["local_user_id"] = local_user.id
                except Exception as sync_err:
                    logger.warning(f"用户同步失败（不影响认证）: {sync_err}")

                # 4. 写入缓存
                token_cache.set(sso_token, user_info)

            # 提取用户信息，优先使用本地用户ID
            # 如果缓存中没有 local_user_id，则重新同步用户
            local_user_id = user_info.get("local_user_id")
            if not local_user_id:
                try:
                    local_user = OAuthService.sync_user_from_oauth(oauth_user_info=user_info)
                    local_user_id = local_user.id
                    # 更新缓存，避免下次重复同步
                    user_info["local_user_id"] = local_user_id
                    token_cache.set(sso_token, user_info)
                except Exception as sync_err:
                    logger.warning(f"重新同步用户失败，使用SSO ID: {sync_err}")
            
            user_id = local_user_id or user_info.get("id") or user_info.get("sub") if not user_info.get("is_admin_token") else user_info.get("id")
            username = (
                user_info.get("username")
                or user_info.get("preferred_username")
                or user_info.get("name")
            )
            email = user_info.get("email", "")
            role = user_info.get("role", "viewer")

            # Attach user info to request state
            request.state.user = {
                "id": str(user_id),
                "username": username,
                "email": email,
                "role": role,
            }

            response = await call_next(request)
            return response

        except HTTPException as e:
            error_type = "invalid_token"
            if e.status_code == 503:
                error_type = "sso_unavailable"
            elif e.status_code == 401:
                # SSO 返回 401 说明 token 过期或无效，统一标记为 token_expired
                # 让前端有机会用 refresh_token 刷新
                error_type = "token_expired"
                # 同时清除本地缓存中的过期 token
                token_cache.invalidate(sso_token)

            return JSONResponse(
                status_code=e.status_code,
                content={
                    "detail": e.detail,
                    "error_type": error_type
                }
            )
        except Exception as e:
            logger.error("认证过程发生错误：%s", str(e))
            return JSONResponse(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                content={
                    "detail": "认证过程发生错误",
                    "error_type": "auth_error"
                }
            )


def require_role(*allowed_roles: str):
    """
    Decorator to check user role.

    Usage:
        @require_role("admin", "annotator")
        async def my_endpoint(request: Request):
            ...

    Args:
        allowed_roles: Tuple of allowed role names

    Returns:
        Decorator function
    """
    def decorator(func):
        async def wrapper(request: Request, *args, **kwargs):
            user = getattr(request.state, "user", None)

            if not user:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="未认证"
                )

            if user["role"] not in allowed_roles:
                raise HTTPException(
                    status_code=status.HTTP_403_FORBIDDEN,
                    detail="权限不足"
                )

            return await func(request, *args, **kwargs)
        return wrapper
    return decorator