from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.models.base import async_session from app.models.user import User from app.services.auth import decode_access_token, user_must_set_password from app.services.credit.query_service import attach_credit_snapshot, get_available_credits security = HTTPBearer(auto_error=False) async def get_db(): async with async_session() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def get_current_user_allow_password_pending( credentials: HTTPAuthorizationCredentials | None = Depends(security), db: AsyncSession = Depends(get_db), ) -> User: if not credentials: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="未登录", ) payload = decode_access_token(credentials.credentials) if not payload: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="登录已过期", ) user_id = payload.get("sub") token_version = payload.get("ver", 0) # Skip captcha tokens if user_id and user_id.startswith("captcha:"): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的凭证", ) result = await db.execute(select(User).where(User.id == user_id).limit(1)) user = result.scalar_one_or_none() if not user or not user.is_active: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="账号不存在或已禁用", ) # 单设备登录校验 — token 版本号不匹配说明已被踢出 if token_version != user.token_version: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="账号已在其他设备登录,请重新登录", ) attach_credit_snapshot(user, await get_available_credits(db, user.id)) return user async def get_current_user( current_user: User = Depends(get_current_user_allow_password_pending), ) -> User: if user_must_set_password(current_user): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={ "code": "PASSWORD_REQUIRED", "message": "请先设置登录密码", }, ) return current_user async def get_optional_current_user( credentials: HTTPAuthorizationCredentials | None = Depends(security), db: AsyncSession = Depends(get_db), ) -> User | None: if not credentials: return None user_id = decode_access_token(credentials.credentials) if not user_id: return None if user_id.startswith("captcha:"): return None result = await db.execute(select(User).where(User.id == user_id).limit(1)) user = result.scalar_one_or_none() if not user or not user.is_active: return None if user_must_set_password(user): return None attach_credit_snapshot(user, await get_available_credits(db, user.id)) return user async def get_admin_user( current_user: User = Depends(get_current_user_allow_password_pending), ) -> User: """验证后台用户权限。 - user_type="admin" 的后台用户即可通过(含非管理员子账号) - 前端通过 allowed_menus 控制非管理员子账号的菜单可见性 - 非后台用户: 403 拒绝 """ if current_user.user_type != "admin": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="需要后台用户权限", ) return current_user async def get_backend_user( current_user: User = Depends(get_current_user_allow_password_pending), ) -> User: """与 get_admin_user 等价: 验证 user_type="admin" 的后台用户。""" if current_user.user_type != "admin": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="需要后台用户权限", ) return current_user