""" 快照构建器 负责从多个数据源聚合数据构建权限快照 """ import asyncio import time from typing import Any from zhenxun.models.ban_console import BanConsole from zhenxun.models.bot_console import BotConsole from zhenxun.models.group_console import GroupConsole from zhenxun.models.level_user import LevelUser from zhenxun.models.plugin_info import PluginInfo from zhenxun.models.user_console import UserConsole from zhenxun.services.log import logger from .models import AuthSnapshot, PluginSnapshot LOG_COMMAND = "auth_snapshot" BUILD_TIMEOUT = 5.0 # 构建超时时间(秒) class SnapshotBuilder: """快照构建器 从多个数据源并行获取数据,聚合成权限快照 """ @classmethod async def build_auth_snapshot( cls, user_id: str, group_id: str | None, bot_id: str, ) -> AuthSnapshot: """构建权限快照 并行获取所有需要的数据,聚合成一个快照对象 参数: user_id: 用户ID group_id: 群组ID(可为None表示私聊) bot_id: Bot ID 返回: AuthSnapshot: 权限快照对象 """ start_time = time.time() try: # 创建所有查询任务 tasks: dict[str, asyncio.Task] = {} # 用户信息 tasks["user"] = asyncio.create_task( cls._get_user(user_id), name="snapshot:user" ) # Ban状态(用户+群组) tasks["ban"] = asyncio.create_task( cls._get_ban_status(user_id, group_id), name="snapshot:ban" ) # 用户权限等级(全局+群组) tasks["level"] = asyncio.create_task( cls._get_user_levels(user_id, group_id), name="snapshot:level" ) # 群组信息(如果有) if group_id: tasks["group"] = asyncio.create_task( cls._get_group(group_id), name="snapshot:group" ) # Bot信息 tasks["bot"] = asyncio.create_task( cls._get_bot(bot_id), name="snapshot:bot" ) # 并行执行所有查询 results: dict[str, Any] = {} try: await asyncio.wait_for( cls._gather_results(tasks, results), timeout=BUILD_TIMEOUT ) except asyncio.TimeoutError: logger.warning( f"构建权限快照超时: user={user_id}, group={group_id}", LOG_COMMAND, ) # 取消未完成的任务 for task in tasks.values(): if not task.done(): task.cancel() # 聚合结果 snapshot = cls._aggregate_results(user_id, group_id, bot_id, results) elapsed = time.time() - start_time if elapsed > 0.5: logger.warning( f"构建权限快照耗时较长: {elapsed:.3f}s, " f"user={user_id}, group={group_id}", LOG_COMMAND, ) return snapshot except Exception as e: logger.error( f"构建权限快照失败: user={user_id}, group={group_id}", LOG_COMMAND, e=e, ) # 返回一个默认快照 return AuthSnapshot(user_id=user_id, group_id=group_id, bot_id=bot_id) @classmethod async def _gather_results( cls, tasks: dict[str, asyncio.Task], results: dict[str, Any] ): """收集所有任务结果 参数: tasks: 任务字典 results: 结果字典(会被修改) """ done, _ = await asyncio.wait(tasks.values(), return_when=asyncio.ALL_COMPLETED) for name, task in tasks.items(): if task in done: try: results[name] = task.result() except Exception as e: logger.warning(f"获取 {name} 数据失败: {e}", LOG_COMMAND) results[name] = None @classmethod async def _get_user(cls, user_id: str) -> UserConsole | None: """获取用户信息""" try: return await UserConsole.get_or_none(user_id=user_id) except Exception as e: logger.warning(f"获取用户信息失败: {user_id}", LOG_COMMAND, e=e) return None @classmethod async def _get_ban_status( cls, user_id: str, group_id: str | None ) -> dict[str, Any]: """获取ban状态 返回: dict: { "user_banned": int, # 0/时间戳/-1 "user_ban_duration": int, "group_banned": int } """ result = { "user_banned": 0, "user_ban_duration": 0, "group_banned": 0, } try: # 获取所有相关的ban记录 ban_records = await BanConsole.is_ban(user_id, group_id) for record in ban_records: if record.user_id and not record.group_id: # 用户级别的ban(全局) if record.duration == -1: result["user_banned"] = -1 result["user_ban_duration"] = -1 else: result["user_banned"] = int(record.ban_time + record.duration) result["user_ban_duration"] = record.duration elif record.user_id and record.group_id: # 用户在特定群组的ban if record.duration == -1: result["user_banned"] = -1 result["user_ban_duration"] = -1 else: result["user_banned"] = int(record.ban_time + record.duration) result["user_ban_duration"] = record.duration elif not record.user_id and record.group_id: # 群组级别的ban if record.duration == -1: result["group_banned"] = -1 else: result["group_banned"] = int(record.ban_time + record.duration) except Exception as e: logger.warning( f"获取ban状态失败: user={user_id}, group={group_id}", LOG_COMMAND, e=e, ) return result @classmethod async def _get_user_levels( cls, user_id: str, group_id: str | None ) -> dict[str, int]: """获取用户权限等级 返回: dict: {"global": int, "group": int} """ result = {"global": 0, "group": 0} try: # 并行查询全局和群组权限 tasks = [] # 全局权限 tasks.append(LevelUser.get_or_none(user_id=user_id, group_id__isnull=True)) # 群组权限 if group_id: tasks.append(LevelUser.get_or_none(user_id=user_id, group_id=group_id)) results = await asyncio.gather(*tasks, return_exceptions=True) # 处理全局权限 if len(results) > 0 and isinstance(results[0], LevelUser): result["global"] = results[0].user_level # 处理群组权限 if len(results) > 1 and isinstance(results[1], LevelUser): result["group"] = results[1].user_level except Exception as e: logger.warning( f"获取用户权限等级失败: user={user_id}, group={group_id}", LOG_COMMAND, e=e, ) return result @classmethod async def _get_group(cls, group_id: str) -> GroupConsole | None: """获取群组信息""" try: return await GroupConsole.get_or_none( group_id=group_id, channel_id__isnull=True ) except Exception as e: logger.warning(f"获取群组信息失败: {group_id}", LOG_COMMAND, e=e) return None @classmethod async def _get_bot(cls, bot_id: str) -> BotConsole | None: """获取Bot信息""" try: return await BotConsole.get_or_none(bot_id=bot_id) except Exception as e: logger.warning(f"获取Bot信息失败: {bot_id}", LOG_COMMAND, e=e) return None @classmethod def _aggregate_results( cls, user_id: str, group_id: str | None, bot_id: str, results: dict[str, Any], ) -> AuthSnapshot: """聚合查询结果为快照 参数: user_id: 用户ID group_id: 群组ID bot_id: Bot ID results: 查询结果字典 返回: AuthSnapshot: 权限快照 """ snapshot = AuthSnapshot( user_id=user_id, group_id=group_id, bot_id=bot_id, ) # 用户信息 if user := results.get("user"): snapshot.user_gold = user.gold # Ban状态 if ban_status := results.get("ban"): snapshot.user_banned = ban_status.get("user_banned", 0) snapshot.user_ban_duration = ban_status.get("user_ban_duration", 0) snapshot.group_banned = ban_status.get("group_banned", 0) # 用户权限等级 if levels := results.get("level"): snapshot.user_level_global = levels.get("global", 0) snapshot.user_level_group = levels.get("group", 0) # 群组信息 if group := results.get("group"): snapshot.group_exists = True snapshot.group_status = group.status snapshot.group_level = group.level snapshot.group_is_super = group.is_super snapshot.group_block_plugins = group.block_plugin or "" snapshot.group_superuser_block_plugins = group.superuser_block_plugin or "" elif group_id: # 有 group_id 但没有群组数据,可能是新群 snapshot.group_exists = False # Bot信息 if bot := results.get("bot"): snapshot.bot_status = bot.status # BotConsole 的 block_plugins 是一个列表 if hasattr(bot, "block_plugins") and bot.block_plugins: if isinstance(bot.block_plugins, list): snapshot.bot_block_plugins = "".join( f"<{p}," for p in bot.block_plugins ) else: snapshot.bot_block_plugins = bot.block_plugins return snapshot @classmethod async def build_plugin_snapshot(cls, module: str) -> PluginSnapshot | None: """构建插件快照 参数: module: 插件模块名 返回: PluginSnapshot | None: 插件快照,不存在时返回None """ try: plugin = await PluginInfo.get_or_none(module=module) if not plugin: return None return PluginSnapshot( module=plugin.module, name=plugin.name, status=plugin.status, block_type=plugin.block_type, plugin_type=plugin.plugin_type, admin_level=plugin.admin_level or 0, cost_gold=plugin.cost_gold, level=plugin.level, limit_superuser=plugin.limit_superuser, ignore_prompt=plugin.ignore_prompt, ) except Exception as e: logger.error(f"构建插件快照失败: {module}", LOG_COMMAND, e=e) return None