zhenxun_bot/zhenxun/builtin_plugins/hooks/auth_snapshot.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from zhenxun.services.cache.runtime_cache import (
    BanMemoryCache,
    BotMemoryCache,
    BotSnapshot,
    GroupMemoryCache,
    GroupSnapshot,
    LevelUserMemoryCache,
    LevelUserSnapshot,
)

from .auth.context import EventContext
from .auth_profile import PluginAuthProfile

if TYPE_CHECKING:
    from nonebot.adapters import Bot


@dataclass(slots=True)
class AuthSnapshot:
    context: EventContext
    plugin: object
    profile: PluginAuthProfile
    bot_data: BotSnapshot | None = None
    group: GroupSnapshot | None = None
    admin_levels: tuple[LevelUserSnapshot | None, LevelUserSnapshot | None] | None = (
        None
    )
    ban_state: bool | None = None
    user_balance_loaded: bool = False
    user_balance: int | None = None
    cache_misses: frozenset[str] = field(default_factory=frozenset)

    @property
    def module(self) -> str:
        return self.profile.module

    @property
    def is_superuser(self) -> bool:
        return self.context.is_superuser

    @property
    def user_id(self) -> str:
        return self.context.user_id

    @property
    def group_id(self) -> str | None:
        return self.context.group_id

    @property
    def channel_id(self) -> str | None:
        return self.context.channel_id

    @property
    def has_ban_cache(self) -> bool:
        return self.ban_state is not None

    @property
    def cache_ready(self) -> bool:
        return not self.cache_misses


async def build_auth_snapshot(
    *,
    context: EventContext,
    plugin: object,
    profile: PluginAuthProfile,
    bot: "Bot",
    skip_ban: bool = False,
    allow_cache_load: bool = False,
) -> AuthSnapshot:
    event_cache = context.event_cache
    entity = context.entity
    cache_misses: set[str] = set()

    bot_data: BotSnapshot | None = None
    if (
        event_cache is not None
        and "bot_data" in event_cache
        and (event_cache.get("bot_cache_ready") or not allow_cache_load)
    ):
        bot_data = event_cache.get("bot_data")
    else:
        bot_data = BotMemoryCache.get_if_ready(bot.self_id)
        if bot_data is None:
            if allow_cache_load:
                bot_data = await BotMemoryCache.get(bot.self_id)
            elif not BotMemoryCache.is_loaded():
                cache_misses.add("bot")
        if event_cache is not None:
            event_cache["bot_data"] = bot_data
            event_cache["bot_cache_ready"] = BotMemoryCache.is_loaded()

    group = None
    if entity.group_id:
        if (
            event_cache is not None
            and "group" in event_cache
            and (event_cache.get("group_cache_ready") or not allow_cache_load)
        ):
            group = event_cache.get("group")
        else:
            group = GroupMemoryCache.get_if_ready(entity.group_id, entity.channel_id)
            if group is None and not GroupMemoryCache.is_loaded():
                cache_misses.add("group")
            elif group is None and allow_cache_load:
                group = await GroupMemoryCache.get(entity.group_id, entity.channel_id)
            if event_cache is not None:
                event_cache["group"] = group
                event_cache["group_cache_ready"] = GroupMemoryCache.is_loaded()

    admin_levels = None
    if profile.need_admin:
        if (
            event_cache is not None
            and "admin_levels" in event_cache
            and (event_cache.get("admin_cache_ready") or not allow_cache_load)
        ):
            admin_levels = event_cache.get("admin_levels")
        else:
            admin_levels = LevelUserMemoryCache.get_levels_if_ready(
                entity.user_id,
                entity.group_id,
            )
            if admin_levels is None:
                if allow_cache_load:
                    admin_levels = await LevelUserMemoryCache.get_levels(
                        entity.user_id,
                        entity.group_id,
                    )
                else:
                    cache_misses.add("admin_levels")
            if event_cache is not None:
                event_cache["admin_levels"] = admin_levels
                event_cache["admin_cache_ready"] = LevelUserMemoryCache.is_loaded()

    ban_state = None
    if not skip_ban:
        if event_cache is not None and "ban_state" in event_cache:
            ban_state = event_cache.get("ban_state")
        elif BanMemoryCache.is_loaded():
            ban_state = BanMemoryCache.is_banned(entity.user_id, entity.group_id)
            if event_cache is not None:
                event_cache["ban_state"] = ban_state
        elif allow_cache_load:
            await BanMemoryCache.ensure_loaded()
            ban_state = BanMemoryCache.is_banned(entity.user_id, entity.group_id)
            if event_cache is not None:
                event_cache["ban_state"] = ban_state
        else:
            cache_misses.add("ban")

    return AuthSnapshot(
        context=context,
        plugin=plugin,
        profile=profile,
        bot_data=bot_data,
        group=group,
        admin_levels=admin_levels,
        ban_state=ban_state,
        cache_misses=frozenset(cache_misses),
    )


async def get_or_build_auth_snapshot(
    *,
    context: EventContext,
    plugin: object,
    profile: PluginAuthProfile,
    bot: "Bot",
    skip_ban: bool = False,
    allow_cache_load: bool = False,
) -> AuthSnapshot:
    event_cache = context.event_cache
    module = profile.module
    if event_cache is not None:
        snapshot_cache = event_cache.setdefault("auth_snapshots", {})
        cached = snapshot_cache.get(module)
        if isinstance(cached, AuthSnapshot):
            if not (allow_cache_load and cached.cache_misses):
                return cached
    snapshot = await build_auth_snapshot(
        context=context,
        plugin=plugin,
        profile=profile,
        bot=bot,
        skip_ban=skip_ban,
        allow_cache_load=allow_cache_load,
    )
    if event_cache is not None:
        event_cache.setdefault("auth_snapshots", {})[module] = snapshot
    return snapshot


__all__ = ["AuthSnapshot", "build_auth_snapshot", "get_or_build_auth_snapshot"]