zhenxun_bot/zhenxun/builtin_plugins/hooks/auth_hook.py

import asyncio
import contextlib
import time

from nonebot import get_driver
from nonebot.adapters import Bot, Event
from nonebot.exception import IgnoredException
from nonebot.matcher import Matcher
from nonebot.message import event_preprocessor, run_postprocessor, run_preprocessor
from nonebot.typing import T_State
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo

from zhenxun.services.cache.runtime_cache import is_cache_ready
from zhenxun.services.log import logger
from zhenxun.services.message_load import is_overloaded
from zhenxun.services.runtime_bootstrap import register_runtime_bootstrap
from zhenxun.utils.utils import get_entity_ids

from .auth.config import LOGGER_COMMAND
from .auth_checker import (
    LimitManager,
    _get_event_cache,
    _get_route_context,
    auth,
    route_precheck,
    start_auth_runtime_tasks,
    stop_auth_runtime_tasks,
)

_SKIP_AUTH_PLUGINS = {"chat_history", "chat_message"}
_BOT_CONNECT_TS: float | None = None
_AUTH_QUEUE_MAXSIZE = 200
_AUTH_QUEUE: asyncio.Queue[tuple[Matcher, Event, Bot, Uninfo, UniMsg | None]] = (
    asyncio.Queue(maxsize=_AUTH_QUEUE_MAXSIZE)
)
_AUTH_QUEUE_STARTED = False
_AUTH_WORKERS: list[asyncio.Task] = []
_LAST_DROP_LOG = 0.0

driver = get_driver()
register_runtime_bootstrap(driver)


@driver.on_bot_connect
async def _mark_bot_connected(bot: Bot):
    del bot
    global _BOT_CONNECT_TS
    _BOT_CONNECT_TS = time.time()


async def _auth_worker(worker_id: int) -> None:
    while True:
        matcher, event, bot, session, message = await _AUTH_QUEUE.get()
        try:
            await auth(
                matcher,
                event,
                bot,
                session,
                message,
                skip_ban=True,
            )
        except IgnoredException:
            pass
        except Exception as exc:
            if not is_overloaded():
                logger.error("async auth failed", LOGGER_COMMAND, e=exc)
        finally:
            _AUTH_QUEUE.task_done()


def _extract_plain_text(message: UniMsg | None, event: Event) -> str:
    if message is not None:
        with contextlib.suppress(Exception):
            return message.extract_plain_text()
    with contextlib.suppress(Exception):
        plain = event.get_plaintext()
        if plain:
            return plain.strip()
    return ""


@driver.on_startup
async def _start_auth_queue():
    global _AUTH_QUEUE_STARTED
    if _AUTH_QUEUE_STARTED:
        return
    _AUTH_QUEUE_STARTED = True
    worker_count = max(1, min(6, _AUTH_QUEUE_MAXSIZE // 50))
    for idx in range(worker_count):
        _AUTH_WORKERS.append(asyncio.create_task(_auth_worker(idx)))
    await start_auth_runtime_tasks()


@driver.on_shutdown
async def _stop_auth_queue():
    global _AUTH_QUEUE_STARTED
    _AUTH_QUEUE_STARTED = False
    workers = _AUTH_WORKERS.copy()
    _AUTH_WORKERS.clear()
    for task in workers:
        task.cancel()
    if workers:
        await asyncio.gather(*workers, return_exceptions=True)

    while not _AUTH_QUEUE.empty():
        with contextlib.suppress(Exception):
            _AUTH_QUEUE.get_nowait()
            _AUTH_QUEUE.task_done()

    await stop_auth_runtime_tasks()


def _skip_auth_for_plugin(matcher: Matcher) -> bool:
    if not matcher.plugin:
        return False
    name = (matcher.plugin.name or "").lower()
    if name in _SKIP_AUTH_PLUGINS:
        return True
    module_name = getattr(matcher.plugin, "module_name", "") or ""
    return "chat_history" in module_name


def _resolve_actor_user_id(event: Event, fallback_user_id: str) -> str:
    """优先使用事件发起者ID，避免 notice 场景 session.user 指向 bot 自身。"""
    event_user_id = getattr(event, "user_id", None)
    if event_user_id is None:
        return fallback_user_id
    event_user_id = str(event_user_id)
    return event_user_id or fallback_user_id


def _resolve_event_group_id(event: Event, fallback_group_id: str | None) -> str | None:
    """notice 场景 session.group 可能缺失，回退到事件上的 group_id。"""
    event_group_id = getattr(event, "group_id", None)
    if event_group_id is None:
        return fallback_group_id
    resolved = str(event_group_id)
    return resolved or fallback_group_id


def _resolve_event_channel_id(
    event: Event, fallback_channel_id: str | None
) -> str | None:
    """频道场景回退到事件上的 channel_id。"""
    event_channel_id = getattr(event, "channel_id", None)
    if event_channel_id is None:
        return fallback_channel_id
    resolved = str(event_channel_id)
    return resolved or fallback_channel_id


@event_preprocessor
async def _drop_message_before_cache_ready(event: Event):
    if event.get_type() != "message":
        return
    if not is_cache_ready():
        raise IgnoredException("cache not ready ignore")
    if _BOT_CONNECT_TS is not None:
        event_ts = getattr(event, "time", None)
        if event_ts is not None and event_ts < _BOT_CONNECT_TS:
            raise IgnoredException("drop backlog message")


@run_preprocessor
async def _auth_preprocessor(
    matcher: Matcher,
    event: Event,
    bot: Bot,
    session: Uninfo,
    state: T_State,
    message: UniMsg | None = None,
):
    if event.get_type() == "message" and not is_cache_ready():
        raise IgnoredException("cache not ready ignore")

    # 提前判断是否跳过权限检查
    if _skip_auth_for_plugin(matcher):
        return

    start_time = time.time()
    entity = state.get("_zx_entity")
    if entity is None:
        entity = get_entity_ids(session)
        entity.user_id = _resolve_actor_user_id(event, entity.user_id)
        entity.group_id = _resolve_event_group_id(event, entity.group_id)
        entity.channel_id = _resolve_event_channel_id(event, entity.channel_id)
        state["_zx_entity"] = entity

    event_cache = state.get("_zx_event_cache")
    if event_cache is None:
        event_cache = _get_event_cache(event, session, entity)
        state["_zx_event_cache"] = event_cache

    text = state.get("_zx_plain_text")
    if text is None:
        text = _extract_plain_text(message, event)
        state["_zx_plain_text"] = text
        if event_cache is not None:
            event_cache["plain_text"] = text

    route_modules = state.get("_zx_route_modules")
    if route_modules is None:
        route_modules = await _get_route_context(text, event_cache)
        state["_zx_route_modules"] = route_modules

    is_superuser = state.get("_zx_is_superuser")
    if is_superuser is None:
        is_superuser = entity.user_id in bot.config.superusers
        state["_zx_is_superuser"] = is_superuser

    if await route_precheck(
        matcher,
        event,
        session,
        message,
        entity=entity,
        event_cache=event_cache,
        text=text,
        route_modules=route_modules,
    ):
        return

    try:
        await auth(
            matcher,
            event,
            bot,
            session,
            message,
            skip_ban=False,
            entity=entity,
            event_cache=event_cache,
            text=text,
            route_modules=route_modules,
            is_superuser=is_superuser,
        )
    except IgnoredException:
        raise
    except Exception as exc:
        logger.error("auth check failed", LOGGER_COMMAND, e=exc)
        raise IgnoredException("auth failed") from exc

    now = time.monotonic()
    last_log = getattr(_auth_preprocessor, "_last_log", 0.0)
    if now - last_log > 1.0 and not is_overloaded():
        setattr(_auth_preprocessor, "_last_log", now)
        logger.debug(
            f"auth check cost: {time.time() - start_time:.3f}s",
            LOGGER_COMMAND,
        )


@run_postprocessor
async def _unblock_after_matcher(matcher: Matcher, session: Uninfo, event: Event):
    user_id = _resolve_actor_user_id(event, session.user.id)
    group_id = _resolve_event_group_id(event, None)
    channel_id = _resolve_event_channel_id(event, None)
    if session.group:
        if session.group.parent:
            group_id = session.group.parent.id
            channel_id = session.group.id
        else:
            group_id = session.group.id
    if user_id and matcher.plugin:
        module = matcher.plugin.name
        LimitManager.unblock(module, user_id, group_id, channel_id)