NeoBot/core/managers/permission_manager.py

"""
权限管理器模块

该模块负责管理用户权限，支持 admin、op、user 三个权限级别。
以 Redis Hash 作为主要数据源，文件仅用作备份和首次数据迁移。
"""
import orjson
import os
from typing import Dict

from ..utils.logger import logger
from ..utils.singleton import Singleton
from .admin_manager import admin_manager
from .redis_manager import redis_manager
from ..permission import Permission


# 用于从字符串名称查找权限对象的字典
_PERMISSIONS: Dict[str, Permission] = {
    p.value: p for p in Permission
}


class PermissionManager(Singleton):
    """
    权限管理器类

    以 Redis Hash 作为权限数据的唯一真实来源，提供高速的读写能力。
    文件 (permissions.json) 仅用于首次启动时的数据迁移和作为灾备。
    """
    _REDIS_KEY = "neobot:permissions"  # 用于存储用户权限的 Redis Hash 键

    def __init__(self):
        """
        初始化权限管理器
        """
        if hasattr(self, '_initialized') and self._initialized:
            return

        # 权限数据文件路径，主要用于备份和首次迁移
        self.data_file = os.path.join(
            os.path.dirname(os.path.abspath(__file__)),
            "..",
            "data",
            "permissions.json"
        )
        
        os.makedirs(os.path.dirname(self.data_file), exist_ok=True)
        logger.info("权限管理器初始化完成")
        super().__init__()

    async def initialize(self):
        """
        异步初始化，检查 Redis 数据，如果为空则尝试从文件迁移
        """
        try:
            if not await redis_manager.redis.exists(self._REDIS_KEY):
                logger.info("Redis 中未找到权限数据，尝试从 permissions.json 文件迁移...")
                await self._migrate_from_file_to_redis()
            else:
                perm_count = await redis_manager.redis.hlen(self._REDIS_KEY)
                logger.info(f"Redis 中已存在权限数据，共 {perm_count} 条。")
        except Exception as e:
            logger.error(f"初始化权限数据时发生错误: {e}")

    async def _migrate_from_file_to_redis(self):
        """
        从 permissions.json 加载权限数据并存入 Redis Hash
        """
        perms_to_migrate = {}
        try:
            if os.path.exists(self.data_file):
                with open(self.data_file, "r", encoding="utf-8") as f:
                    data = orjson.loads(f.read())
                    perms_to_migrate = data.get("users", {})
            
            if perms_to_migrate:
                # 使用 pipeline 批量写入，提高效率
                async with redis_manager.redis.pipeline(transaction=True) as pipe:
                    for user_id, level_name in perms_to_migrate.items():
                        pipe.hset(self._REDIS_KEY, user_id, level_name)
                    await pipe.execute()
                logger.success(f"成功从文件迁移 {len(perms_to_migrate)} 条权限数据到 Redis。")
            else:
                logger.info("permissions.json 文件为空或不存在，无需迁移。")

        except ValueError as e:
            logger.error(f"解析 permissions.json 失败，无法迁移: {e}")
        except Exception as e:
            logger.error(f"迁移权限数据到 Redis 失败: {e}")

    async def _save_to_file_backup(self):
        """
        将 Redis 中的权限数据完整备份到 permissions.json
        """
        try:
            all_perms = await redis_manager.redis.hgetall(self._REDIS_KEY)
            # Redis 返回的是 bytes，需要解码
            users_data = {k.decode('utf-8'): v.decode('utf-8') for k, v in all_perms.items()}
            with open(self.data_file, "w", encoding="utf-8") as f:
                f.write(orjson.dumps({"users": users_data}, indent=2, ensure_ascii=False).decode('utf-8'))
            logger.debug(f"权限数据已备份到 {self.data_file}")
        except Exception as e:
            logger.error(f"备份权限数据到 permissions.json 失败: {e}")

    async def get_user_permission(self, user_id: int) -> Permission:
        """
        获取指定用户的权限对象

        优先检查是否为机器人管理员，然后从 Redis 查询。
        """
        if await admin_manager.is_admin(user_id):
            return Permission.ADMIN

        try:
            level_name_bytes = await redis_manager.redis.hget(self._REDIS_KEY, str(user_id))
            if level_name_bytes:
                level_name = level_name_bytes.decode('utf-8')
                return _PERMISSIONS.get(level_name, Permission.USER)
        except Exception as e:
            logger.error(f"从 Redis 获取用户 {user_id} 权限失败: {e}")
        
        return Permission.USER

    async def set_user_permission(self, user_id: int, permission: Permission) -> None:
        """
        在 Redis 中设置指定用户的权限级别，并更新文件备份
        """
        if not isinstance(permission, Permission):
            raise ValueError(f"无效的权限对象: {permission}")

        try:
            await redis_manager.redis.hset(self._REDIS_KEY, str(user_id), permission.value)
            await self._save_to_file_backup()
            logger.info(f"已在 Redis 中设置用户 {user_id} 的权限为 {permission.value}")
        except Exception as e:
            logger.error(f"在 Redis 中设置用户 {user_id} 权限失败: {e}")

    async def remove_user(self, user_id: int) -> None:
        """
        从 Redis 中移除指定用户的权限设置，并更新文件备份
        """
        try:
            if await redis_manager.redis.hdel(self._REDIS_KEY, str(user_id)):
                await self._save_to_file_backup()
                logger.info(f"已从 Redis 中移除用户 {user_id} 的权限设置")
        except Exception as e:
            logger.error(f"从 Redis 移除用户 {user_id} 权限失败: {e}")

    async def check_permission(self, user_id: int, required_permission: Permission) -> bool:
        """
        检查用户是否具有指定权限级别
        """
        user_permission = await self.get_user_permission(user_id)
        
        # 增强类型检查，防止将property对象等错误类型传递进来
        if not isinstance(required_permission, Permission):
            logger.error(f"权限检查失败：required_permission 不是 Permission 枚举类型，而是 {type(required_permission).__name__}")
            return False
            
        return user_permission >= required_permission

    async def get_all_user_permissions(self) -> Dict[str, str]:
        """
        获取所有已配置的用户权限（合并 Redis 和 AdminManager）
        """
        permissions = {}
        try:
            # 从 Redis 获取基础权限
            all_perms = await redis_manager.redis.hgetall(self._REDIS_KEY)
            permissions = {k.decode('utf-8'): v.decode('utf-8') for k, v in all_perms.items()}
        except Exception as e:
            logger.error(f"从 Redis 获取所有权限失败: {e}")

        # 合并 AdminManager 中的管理员，ADMIN 权限覆盖一切
        try:
            admins = await admin_manager.get_all_admins()
            for admin_id in admins:
                permissions[str(admin_id)] = Permission.ADMIN.value
        except Exception as e:
            logger.error(f"获取管理员列表以合并权限时失败: {e}")
            
        return permissions

    async def clear_all(self) -> None:
        """
        清空 Redis 中的所有权限设置，并更新备份文件
        """
        try:
            await redis_manager.redis.delete(self._REDIS_KEY)
            await self._save_to_file_backup()
            logger.info("已清空 Redis 中的所有权限设置")
        except Exception as e:
            logger.error(f"清空 Redis 权限数据失败: {e}")


def require_admin(func):
    """
    一个装饰器，用于限制命令只能由管理员执行。
    """
    from functools import wraps
    from models.events.message import MessageEvent
    from core.managers import permission_manager

    @wraps(func)
    async def wrapper(event: MessageEvent, *args, **kwargs):
        user_id = event.user_id
        if await permission_manager.check_permission(user_id, Permission.ADMIN):
            return await func(event, *args, **kwargs)
        else:
            # 假设 event 对象有 reply 方法
            if hasattr(event, "reply"):
                await event.reply("抱歉，您没有权限执行此命令。")
            return None
    return wrapper