chat-ai-api/app/dependencies/auth.py

from fastapi import HTTPException, Query, Depends
from fastapi.security import OAuth2PasswordBearer
from typing import Optional, Annotated
from pydantic import BaseModel
import jwt

from ..db.redis_client import get_app_user
from ..config.config import Config

config = Config()
oauth2_scheme_optional = OAuth2PasswordBearer(tokenUrl="/users/token", auto_error=False)


class UserContext(BaseModel):
    user_id: str
    app_name: str


async def resolve_user_id(
    jwt_token: Annotated[Optional[str], Depends(oauth2_scheme_optional)] = None,
    source: Optional[str] = Query(default=None),
    token: Optional[str] = Query(default=None),
) -> UserContext:
    if source == "app" and token:
        app_user = get_app_user(token)
        if not app_user:
            raise HTTPException(status_code=401, detail="无效的 App token")

        user_id = app_user.get("userId")
        app_name = app_user.get("appName")
        if not user_id or not app_name:
            raise HTTPException(status_code=401, detail="无效的 App token")

        return UserContext(
            user_id=user_id,
            app_name=app_name,
        )

    if not jwt_token:
        raise HTTPException(status_code=401, detail="未提供认证令牌")

    try:
        payload = jwt.decode(jwt_token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        user_id = payload.get("userId")
        if not user_id:
            raise HTTPException(status_code=401, detail="无效的令牌")

        return UserContext(
            user_id=user_id,
            app_name="com.yunxingshnegtai",
        )
    except jwt.PyJWTError:
        raise HTTPException(status_code=401, detail="无效的令牌")