FastapiAdmin/backend/app/core/dependencies.py

# -*- coding: utf-8 -*-

import json
from redis.asyncio.client import Redis
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
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.api.v1.module_system.user.model import UserModel
from app.api.v1.module_system.role.model import RoleModel
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,
    db: AsyncSession = Depends(db_getter),
    redis: Redis = Depends(redis_getter),
    token: str = Depends(OAuth2Schema),
) -> AuthSchema:
    """获取当前用户
    
    参数:
    - request (Request): 请求对象
    - db (AsyncSession): 数据库会话
    - redis (Redis): Redis连接
    - token (str): 访问令牌
    
    返回:
    - AuthSchema: 认证信息模型
    """
    if not token:
        raise CustomException(msg="认证已失效", code=10401, status_code=401)
    
    # 处理Bearer token
    if token.startswith('Bearer'):
        token = token.split(' ')[1]

    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, check_data_scope=False)
    username = user_info.get("user_name")
    if not username:
        raise CustomException(msg="认证已失效", code=10401, status_code=401)
    # 获取用户信息，使用深层预加载确保RoleModel.creator被正确加载
    user = await UserCRUD(auth).get_by_username_crud(
        username=username, 
        preload=[
            "dept", 
            selectinload(UserModel.roles).selectinload(RoleModel.creator),
            "positions", 
            "creator"
        ]
    )
    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 and role.status]
    if hasattr(user, 'positions'):
        user.positions = [pos for pos in user.positions if pos and 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 = permissions or []
        self.check_data_scope = check_data_scope

    async def __call__(self, auth: AuthSchema = Depends(get_current_user)) -> AuthSchema:
        """
        调用权限验证
        
        参数:
        - auth (AuthSchema): 认证信息对象。
        
        返回:
        - AuthSchema: 认证信息对象。
        """
        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 "*" in self.permissions or "*:*:*" in 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 role.status and menu.permission and menu.status
        }

        # 权限验证 - 满足任一权限即可
        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