zhenxun_bot/zhenxun/builtin_plugins/hooks/auth/auth_limit.py

import asyncio
import time
from typing import Any, ClassVar

import nonebot
from nonebot_plugin_uninfo import Uninfo
from pydantic import BaseModel

from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.plugin_limit import PluginLimit
from zhenxun.services.cache.runtime_cache import (
    PluginLimitMemoryCache,
    PluginLimitSnapshot,
)
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
from zhenxun.services.log import logger
from zhenxun.utils.enum import LimitWatchType, PluginLimitType
from zhenxun.utils.limiters import CountLimiter, FreqLimiter, UserBlockLimiter
from zhenxun.utils.manager.priority_manager import PriorityLifecycle
from zhenxun.utils.message import MessageUtils
from zhenxun.utils.time_utils import TimeUtils
from zhenxun.utils.utils import EntityIDs, get_entity_ids

from .config import LOGGER_COMMAND, WARNING_THRESHOLD
from .context import PermissionContext
from .exception import SkipPluginException

driver = nonebot.get_driver()

_LIMIT_NOTICE_CD = 2
_LIMIT_NOTICE_LIMITER = FreqLimiter(_LIMIT_NOTICE_CD)
_LIMIT_NOTICE_TASKS: set[asyncio.Task] = set()


@PriorityLifecycle.on_startup(priority=7)
async def _():
    """初始化限制"""
    await LimitManager.init_limit()


class Limit(BaseModel):
    limit: PluginLimit | PluginLimitSnapshot
    limiter: FreqLimiter | UserBlockLimiter | CountLimiter

    class Config:
        arbitrary_types_allowed = True


def _limit_notice_key(
    limit: PluginLimit | PluginLimitSnapshot,
    user_id: str,
    group_id: str | None,
    channel_id: str | None,
) -> str:
    key = user_id
    if group_id and limit.watch_type == LimitWatchType.GROUP:
        key = channel_id or group_id
    return f"{limit.module}:{limit.limit_type}:{key}"


def _send_limit_notice(message: str, format_kwargs: dict[str, Any], key: str) -> None:
    if not _LIMIT_NOTICE_LIMITER.check(key):
        return
    _LIMIT_NOTICE_LIMITER.start_cd(key)

    async def _send():
        try:
            await MessageUtils.build_message(message, format_args=format_kwargs).send()
        except Exception as exc:
            logger.error("limit notice send failed", LOGGER_COMMAND, e=exc)

    task = asyncio.create_task(_send())
    _LIMIT_NOTICE_TASKS.add(task)
    task.add_done_callback(_LIMIT_NOTICE_TASKS.discard)


class LimitManager:
    add_module: ClassVar[list] = []
    last_update_time: ClassVar[float] = 0
    update_interval: ClassVar[float] = 6000  # 1小时更新一次
    is_updating: ClassVar[bool] = False  # 防止并发更新

    cd_limit: ClassVar[dict[str, Limit]] = {}
    block_limit: ClassVar[dict[str, Limit]] = {}
    count_limit: ClassVar[dict[str, Limit]] = {}

    # 模块限制缓存，避免频繁查询数据库
    module_limit_cache: ClassVar[
        dict[str, tuple[float, list[PluginLimitSnapshot], bool]]
    ] = {}
    module_cache_ttl: ClassVar[float] = 60  # 模块缓存有效期（秒）
    module_cache_error_ttl: ClassVar[float] = 5  # 超时缓存有效期（秒）

    @classmethod
    async def init_limit(cls):
        """初始化限制"""
        cls.last_update_time = time.time()
        try:
            await asyncio.wait_for(cls.update_limits(), timeout=DB_TIMEOUT_SECONDS * 2)
        except asyncio.TimeoutError:
            logger.error("初始化限制超时", LOGGER_COMMAND)

    @classmethod
    async def update_limits(cls):
        """更新限制信息"""
        # 防止并发更新
        if cls.is_updating:
            return

        cls.is_updating = True
        try:
            start_time = time.time()
            await PluginLimitMemoryCache.ensure_loaded()
            limit_list = PluginLimitMemoryCache.get_all_limits()

            # 清空旧数据
            cls.add_module = []
            cls.cd_limit = {}
            cls.block_limit = {}
            cls.count_limit = {}
            cls.module_limit_cache.clear()
            # 添加新数据
            for limit in limit_list:
                cls.add_limit(limit)

            cls.last_update_time = time.time()
            elapsed = time.time() - start_time
            if elapsed > WARNING_THRESHOLD:  # 记录耗时超过500ms的更新
                logger.warning(f"更新限制信息耗时: {elapsed:.3f}s", LOGGER_COMMAND)
        finally:
            cls.is_updating = False

    @classmethod
    def add_limit(cls, limit: PluginLimit | PluginLimitSnapshot):
        """添加限制

        参数:
            limit: PluginLimit
        """
        if limit.module not in cls.add_module:
            cls.add_module.append(limit.module)
        if limit.limit_type == PluginLimitType.BLOCK:
            cls.block_limit[limit.module] = Limit(
                limit=limit, limiter=UserBlockLimiter()
            )
        elif limit.limit_type == PluginLimitType.CD:
            cd_value = int(limit.cd or 0)
            cls.cd_limit[limit.module] = Limit(
                limit=limit, limiter=FreqLimiter(cd_value)
            )
        elif limit.limit_type == PluginLimitType.COUNT:
            max_count = int(limit.max_count or 0)
            if max_count <= 0:
                return
            cls.count_limit[limit.module] = Limit(
                limit=limit, limiter=CountLimiter(max_count)
            )

    @classmethod
    def unblock(
        cls, module: str, user_id: str, group_id: str | None, channel_id: str | None
    ):
        """解除插件block

        参数:
            module: 模块名
            user_id: 用户id
            group_id: 群组id
            channel_id: 频道id
        """
        if limit_model := cls.block_limit.get(module):
            limit = limit_model.limit
            limiter: UserBlockLimiter = limit_model.limiter  # type: ignore
            key_type = user_id
            if group_id and limit.watch_type == LimitWatchType.GROUP:
                key_type = channel_id or group_id
            logger.debug(
                f"解除对象: {key_type} 的block限制",
                LOGGER_COMMAND,
                session=user_id,
                group_id=group_id,
            )
            limiter.set_false(key_type)

    @classmethod
    async def get_module_limits(cls, module: str) -> list[PluginLimitSnapshot]:
        """获取模块的限制信息，使用缓存减少数据库查询

        参数:
            module: 模块名

        返回:
            list[PluginLimit]: 限制列表
        """
        current_time = time.time()

        # 检查缓存
        if module in cls.module_limit_cache:
            cache_time, limits, is_error = cls.module_limit_cache[module]
            ttl = cls.module_cache_error_ttl if is_error else cls.module_cache_ttl
            if current_time - cache_time < ttl:
                return limits

        # 缓存不存在或已过期，从内存缓存获取
        try:
            await PluginLimitMemoryCache.ensure_loaded()
            limits = await PluginLimitMemoryCache.get_limits(module)
            cls.module_limit_cache[module] = (current_time, limits, False)
            return limits
        except Exception as exc:
            logger.error(f"get module limits failed: {module}", LOGGER_COMMAND, e=exc)
            cls.module_limit_cache[module] = (current_time, [], True)
            return []

    @classmethod
    async def check(
        cls,
        module: str,
        user_id: str,
        group_id: str | None,
        channel_id: str | None,
    ):
        """检测限制

        参数:
            module: 模块名
            user_id: 用户id
            group_id: 群组id
            channel_id: 频道id

        异常:
            IgnoredException: IgnoredException
        """
        start_time = time.time()

        # 定期更新全局限制信息
        if (
            time.time() - cls.last_update_time > cls.update_interval
            and not cls.is_updating
        ):
            # 使用异步任务更新，避免阻塞当前请求
            asyncio.create_task(cls.update_limits())  # noqa: RUF006

        # 如果模块不在已加载列表中，只加载该模块的限制
        if module not in cls.add_module:
            limits = await cls.get_module_limits(module)
            for limit in limits:
                cls.add_limit(limit)

        # 检查各种限制
        try:
            if limit_model := cls.cd_limit.get(module):
                await cls.__check(limit_model, user_id, group_id, channel_id)
            if limit_model := cls.block_limit.get(module):
                await cls.__check(limit_model, user_id, group_id, channel_id)
            if limit_model := cls.count_limit.get(module):
                await cls.__check(limit_model, user_id, group_id, channel_id)
        finally:
            # 记录总执行时间
            elapsed = time.time() - start_time
            if elapsed > WARNING_THRESHOLD:  # 记录耗时超过500ms的检查
                logger.warning(
                    f"限制检查耗时: {elapsed:.3f}s, 模块: {module}",
                    LOGGER_COMMAND,
                    session=user_id,
                    group_id=group_id,
                )

    @classmethod
    async def __check(
        cls,
        limit_model: Limit | None,
        user_id: str,
        group_id: str | None,
        channel_id: str | None,
    ):
        """检测限制

        参数:
            limit_model: Limit
            user_id: 用户id
            group_id: 群组id
            channel_id: 频道id

        异常:
            IgnoredException: IgnoredException
        """
        if not limit_model:
            return
        limit = limit_model.limit
        limiter = limit_model.limiter
        is_limit = (
            LimitWatchType.ALL
            or (group_id and limit.watch_type == LimitWatchType.GROUP)
            or (not group_id and limit.watch_type == LimitWatchType.USER)
        )
        key_type = user_id
        if group_id and limit.watch_type == LimitWatchType.GROUP:
            key_type = channel_id or group_id
        if is_limit and not limiter.check(key_type):
            if limit.result:
                format_kwargs = {}
                if isinstance(limiter, FreqLimiter):
                    left_time = limiter.left_time(key_type)
                    cd_str = TimeUtils.format_duration(left_time)
                    format_kwargs = {"cd": cd_str}
                notice_key = _limit_notice_key(limit, user_id, group_id, channel_id)
                _send_limit_notice(limit.result, format_kwargs, notice_key)
            raise SkipPluginException(
                f"{limit.module}({limit.limit_type}) 正在限制中..."
            )
        else:
            logger.debug(
                f"开始进行限制 {limit.module}({limit.limit_type})...",
                LOGGER_COMMAND,
                session=user_id,
                group_id=group_id,
            )
            if isinstance(limiter, FreqLimiter):
                limiter.start_cd(key_type)
            if isinstance(limiter, UserBlockLimiter):
                limiter.set_true(key_type)
            if isinstance(limiter, CountLimiter):
                limiter.increase(key_type)


async def auth_limit(
    plugin: PluginInfo,
    session: Uninfo,
    *,
    context: PermissionContext | None = None,
    entity: EntityIDs | None = None,
):
    """插件限制

    参数:
        plugin: PluginInfo
        session: Uninfo
    """
    if context is not None:
        entity = context.entity
    if entity is None:
        entity = get_entity_ids(session)
    try:
        await asyncio.wait_for(
            LimitManager.check(
                plugin.module, entity.user_id, entity.group_id, entity.channel_id
            ),
            timeout=DB_TIMEOUT_SECONDS * 2,  # 给予更长的超时时间
        )
    except asyncio.TimeoutError:
        logger.error(f"检查插件限制超时: {plugin.module}", LOGGER_COMMAND)
        # 超时时不抛出异常，允许继续执行