# -*- coding: utf-8 -*- import json from redis.asyncio.client import Redis from sqlalchemy.ext.asyncio import AsyncSession from typing import AsyncGenerator, Optional from fastapi import Depends, Request from motor.motor_asyncio import AsyncIOMotorDatabase from fastapi import Depends from app.api.v1.module_system.user.schema import UserOutSchema from app.common.enums import RedisInitKeyConfig from app.core.exceptions import CustomException from app.core.database import session_connect from app.core.security import OAuth2Schema, decode_access_token from app.core.logger import logger from app.core.redis_crud import RedisCURD from app.api.v1.module_system.user.crud import UserCRUD from app.api.v1.module_system.auth.schema import AuthSchema async def db_getter() -> AsyncGenerator[AsyncSession, None]: """获取数据库会话连接 返回: - AsyncSession: 数据库会话连接 """ async with session_connect() as session: async with session.begin(): yield session async def redis_getter(request: Request) -> Redis: """获取Redis连接 参数: - request (Request): 请求对象 返回: - Redis: Redis连接 """ return request.app.state.redis async def mongo_getter(request: Request) -> AsyncIOMotorDatabase: """获取MongoDB连接 参数: - request (Request): 请求对象 返回: - AsyncIOMotorDatabase: MongoDB连接 """ return request.app.state.mongo async def get_current_user( request: Request, token: str = Depends(OAuth2Schema), redis: Redis = Depends(redis_getter), db: AsyncSession = Depends(db_getter) ) -> AuthSchema: """获取并验证当前用户信息 参数: - request (Request): 请求对象。 - token (str): 认证token。 - redis (Redis): Redis连接。 - db (AsyncSession): 数据库会话连接。 返回: - AuthSchema: 包含用户信息的认证对象。 异常: - CustomException: 认证失败时抛出异常。 """ # 处理Bearer token if token.startswith('Bearer'): token = token.split(' ')[1] # 解析token payload = decode_access_token(token) if not payload or not hasattr(payload, 'is_refresh') or payload.is_refresh: raise CustomException(msg="非法凭证", code=10401, status_code=401) online_user_info = payload.sub # 从Redis中获取用户信息 user_info = json.loads(online_user_info) # 确保是字典类型 session_id = user_info.get("session_id") if not session_id: raise CustomException(msg="认证已失效", code=10401, status_code=401) # 检查用户是否在线 online_ok = await RedisCURD(redis).exists(key=f'{RedisInitKeyConfig.ACCESS_TOKEN.key}:{session_id}') if not online_ok: raise CustomException(msg="认证已失效", code=10401, status_code=401) auth = AuthSchema(db=db) username = user_info.get("user_name") if not username: raise CustomException(msg="认证已失效", code=10401, status_code=401) # 获取用户信息 user = await UserCRUD(auth).get_by_username_crud(username=username) if not user: raise CustomException(msg="用户不存在", code=10401, status_code=401) if not user.status: raise CustomException(msg="用户已被停用", code=10401, status_code=401) # 设置请求上下文 request.scope["user_id"] = user.id request.scope["user_username"] = user.username # 过滤可用的角色和职位 if hasattr(user, 'roles'): user.roles = [role for role in user.roles if role.status] if hasattr(user, 'positions'): user.positions = [pos for pos in user.positions if pos.status] auth.user = UserOutSchema.model_validate(user) return auth class AuthPermission: """权限验证类""" def __init__(self, permissions: Optional[list[str]] = None, check_data_scope: bool = True) -> None: """ 初始化权限验证 参数: - permissions (Optional[list[str]]): 权限标识列表。 - check_data_scope (bool): 是否启用严格模式校验。 """ self.permissions = set(permissions) if permissions else None self.check_data_scope = check_data_scope async def __call__(self, auth: AuthSchema = Depends(get_current_user)) -> AuthSchema: """ 执行权限验证 参数: - request (Request): 请求对象。 - auth (AuthSchema): 认证信息。 返回: - AuthSchema: 认证对象 异常: - CustomException: 权限验证失败时抛出异常。 """ auth.check_data_scope = self.check_data_scope # 超级管理员直接通过 if auth.user and auth.user.is_superuser: return auth # 无需验证权限 if not self.permissions: return auth # 超级管理员权限标识 if {"*:*:*"} <= self.permissions: return auth # 检查用户是否有角色 if not auth.user or not auth.user.roles: raise CustomException(msg="无权限操作", code=10403, status_code=403) # 获取用户权限集合 user_permissions = { menu.permission for role in auth.user.roles for menu in role.menus if menu.permission and menu.status } # 权限验证 if self.check_data_scope: # 严格模式:要求所有权限都满足 if not all(perm in user_permissions for perm in self.permissions): logger.error(f"用户缺少所需的权限: {self.permissions}") raise CustomException(msg="无权限操作", code=10403, status_code=403) else: # 非严格模式:满足任一权限即可 if not any(perm in user_permissions for perm in self.permissions): logger.error(f"用户缺少任何所需的权限: {self.permissions}") raise CustomException(msg="无权限操作", code=10403, status_code=403) return auth