import hashlib import re from fastapi import Depends, Header, HTTPException, status from fastapi.security import OAuth2PasswordBearer from jose import JWTError from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from yuxi.storage.postgres.manager import pg_manager from yuxi.storage.postgres.models_business import APIKey, User from yuxi.utils.datetime_utils import utc_now_naive from server.utils.auth_utils import AuthUtils # 定义OAuth2密码承载器,指定token URL oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/token", auto_error=False) # 公开路径列表,无需登录即可访问 PUBLIC_PATHS = [ r"^/api/auth/token$", # 登录 r"^/api/auth/check-first-run$", # 检查是否首次运行 r"^/api/auth/initialize$", # 初始化系统 r"^/api$", # Health Check r"^/api/system/health$", # Health Check r"^/api/system/info$", # 获取系统信息配置 r"^/api/webhook/.*$", # Webhook 回调(使用渠道签名验证,不走 JWT/API Key) r"^/ws/chat/.*$", # WebSocket 通过查询参数鉴权 ] def _is_channel_user(user: User) -> bool: return bool(user.source and user.source.startswith("channel:")) # 获取数据库会话(异步版本) async def get_db(): async with pg_manager.get_async_session_context() as db: yield db async def _verify_api_key(key: str, db: AsyncSession) -> tuple[User | None, APIKey | None]: """验证 API Key 并返回关联用户和 APIKey 对象""" key_hash = hashlib.sha256(key.encode()).hexdigest() result = await db.execute(select(APIKey).filter(APIKey.key_hash == key_hash)) api_key = result.scalar_one_or_none() if api_key is None: return None, None if not api_key.is_enabled: return None, None if api_key.expires_at and utc_now_naive() > api_key.expires_at: return None, None if api_key.user_id: result = await db.execute(select(User).filter(User.id == api_key.user_id)) user = result.scalar_one_or_none() if user and not user.is_deleted and not _is_channel_user(user): return user, api_key if api_key.department_id: result = await db.execute( select(User).filter(User.department_id == api_key.department_id, User.role.in_(["admin", "superadmin"])) ) user = result.scalar_one_or_none() if user and not user.is_deleted and not _is_channel_user(user): return user, api_key result = await db.execute(select(User).filter(User.role == "superadmin", User.is_deleted == 0).limit(1)) user = result.scalar_one_or_none() if user and not _is_channel_user(user): return user, api_key return None, None # 获取当前用户(异步版本) async def get_current_user( authorization: str | None = Header(None), db: AsyncSession = Depends(get_db), ): credentials_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的凭证", headers={"WWW-Authenticate": "Bearer"}, ) if authorization is None: return None if not authorization.startswith("Bearer "): return None token = authorization.split("Bearer ")[1] if not token: return None # 根据 token 前缀判断认证方式 if token.startswith("yxkey_"): # API Key 认证 user, api_key_obj = await _verify_api_key(token, db) if user is not None and api_key_obj is not None: api_key_obj.last_used_at = utc_now_naive() await db.commit() return user # JWT Token 认证 try: payload = AuthUtils.verify_access_token(token) user_id = payload.get("sub") if user_id is None: raise credentials_exception except JWTError: raise credentials_exception except ValueError as e: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=str(e), headers={"WWW-Authenticate": "Bearer"}, ) result = await db.execute(select(User).filter(User.id == int(user_id), User.is_deleted == 0)) user = result.scalar_one_or_none() if user is None: raise credentials_exception if user.is_login_locked(): raise HTTPException( status_code=status.HTTP_423_LOCKED, detail="登录被锁定,请稍后重试", headers={"X-Lock-Remaining": str(user.get_remaining_lock_time())}, ) if _is_channel_user(user): raise credentials_exception return user # 获取已登录用户(抛出401如果未登录) async def get_required_user(user: User | None = Depends(get_current_user)): if user is None: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="请登录后再访问", headers={"WWW-Authenticate": "Bearer"}, ) if not user.department_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="当前用户未绑定部门", ) return user # 获取管理员用户 async def get_admin_user(current_user: User = Depends(get_required_user)): if current_user.role not in ["admin", "superadmin"]: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="需要管理员权限", ) return current_user # 获取超级管理员用户 async def get_superadmin_user(current_user: User = Depends(get_required_user)): if current_user.role != "superadmin": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="需要超级管理员权限", ) return current_user # 检查路径是否为公开路径 def is_public_path(path: str) -> bool: path = path.rstrip("/") # 去除尾部斜杠以便于匹配 for pattern in PUBLIC_PATHS: if re.match(pattern, path): return True return False