zhenxun_bot/zhenxun/builtin_plugins/hooks/auth_pipeline.py

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
import time
from typing import TYPE_CHECKING, Any

from nonebot.adapters import Bot, Event
from nonebot.matcher import Matcher
from nonebot_plugin_uninfo import Uninfo

from zhenxun.services.message_load import is_db_unhealthy
from zhenxun.utils.utils import EntityIDs

from .auth.context import (
    EventContext,
    PermissionSideEffectCache,
    set_route_modules,
)
from .auth.exception import PermissionExemption, SkipPluginException
from .auth_policy import (
    action_from_snapshot,
    principal_from_snapshot,
    raise_for_policy,
    resource_from_snapshot,
)
from .auth_types import AuthLaneContext, AuthPolicyFlags, AuthPreparation

if TYPE_CHECKING:
    from .auth_side_effect import SideEffectCommit
    from .auth_trace import HookTraceRecorder


def _require(value: Any, name: str):
    if value is None:
        raise RuntimeError(f"AuthPipelineContext.{name} is required")
    return value


def _prep(ctx: AuthPipelineContext) -> AuthPreparation:
    return _require(ctx.prep, "prep")


def _recorder(ctx: AuthPipelineContext) -> "HookTraceRecorder":
    return _require(ctx.hook_recorder, "hook_recorder")


def _side_effect_commit(ctx: AuthPipelineContext) -> "SideEffectCommit":
    return _require(ctx.side_effect_commit, "side_effect_commit")


def _side_effect_cache(ctx: AuthPipelineContext) -> PermissionSideEffectCache:
    return _require(ctx.side_effect_cache, "side_effect_cache")


def _entity(ctx: AuthPipelineContext) -> EntityIDs:
    return _require(ctx.entity, "entity")


def _lane_context(ctx: AuthPipelineContext) -> AuthLaneContext:
    return _require(ctx.lane_context, "lane_context")


PipelineHandler = Callable[["AuthPipelineContext"], Awaitable[None]]


@dataclass(slots=True)
class AuthPipelineStage:
    name: str
    handler: PipelineHandler


@dataclass(slots=True)
class AuthPipelineContext:
    matcher: Matcher
    event: Event
    bot: Bot
    session: Uninfo
    event_context: EventContext
    skip_ban: bool = False
    state: dict | None = None
    start_time: float = field(default_factory=time.time)
    module: str = ""
    entity: EntityIDs | None = None
    event_cache: dict | None = None
    text: str = ""
    route_modules: set[str] | None = None
    is_command_matcher: bool = False
    lane_context: AuthLaneContext | None = None
    side_effect_cache: PermissionSideEffectCache | None = None
    side_effect_commit: "SideEffectCommit | None" = None
    side_effect_lock: asyncio.Lock | None = None
    entered_side_effect_lock: bool = False
    auth_result_cache: dict | None = None
    hook_recorder: "HookTraceRecorder | None" = None
    prep: AuthPreparation | None = None
    flags: AuthPolicyFlags | None = None
    cost_gold: int = 0
    hooks_time: float = 0.0
    ignore_flag: bool = False
    auth_allowed: bool | None = None
    decision_effect: str | None = None
    decision_reason: str | None = None
    stopped: bool = False
    stage_timings: dict[str, float] = field(default_factory=dict)

    def stop(
        self,
        *,
        allowed: bool,
        effect: str,
        reason: str,
    ) -> None:
        self.auth_allowed = allowed
        self.decision_effect = effect
        self.decision_reason = reason
        self.stopped = True


class AuthPipeline:
    def __init__(self, stages: list[AuthPipelineStage]) -> None:
        self._stages = tuple(stages)

    async def run(self, context: AuthPipelineContext) -> None:
        for stage in self._stages:
            started = time.perf_counter()
            await stage.handler(context)
            context.stage_timings[stage.name] = (time.perf_counter() - started) * 1000
            if context.stopped:
                break


@dataclass(slots=True)
class AuthPipelineDependencies:
    route_modules_with_commands: set[str]
    get_route_context: Callable[[str, dict | None], Awaitable[set[str]]]
    is_hidden_plugin: Callable[[Matcher], bool]
    is_command_matcher_class: Callable[[type[Matcher]], bool]
    matcher_has_alconna_shortcuts: Callable[[type[Matcher]], bool]
    prepare_auth_state_with_fallback: Callable[..., Awaitable[Any]]
    prepare_auth_state: Callable[..., Awaitable[Any]]
    policy_decision_point: Any
    policy_skip_message: Callable[[str], str]
    legacy_pure_auth_fallback: Callable[..., Awaitable[None]]
    check_ban_from_snapshot: Callable[..., Awaitable[None]]
    resolve_cost_gold: Callable[..., Awaitable[int]]
    run_auth_hooks: Callable[..., Awaitable[float]]
    bot_filter: Callable[..., None]
    reserve_gold: Callable[..., Awaitable[Any]]
    insufficient_gold_error: type[Exception]
    logger: Any
    log_command: str


def apply_policy_precheck(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> AuthPolicyFlags:
    prep = _prep(ctx)
    hook_recorder = _recorder(ctx)
    flags = AuthPolicyFlags()
    snapshot = prep.snapshot
    decision = deps.policy_decision_point.decide(
        principal_from_snapshot(snapshot),
        action_from_snapshot(snapshot),
        resource_from_snapshot(snapshot),
        prep.policy_context,
    )
    if decision.deferred:
        hook_recorder.set("auth_core", f"policy:{decision.reason}")
    if decision.denied:
        raise_for_policy(decision, deps.policy_skip_message(decision.reason))
    if decision.allowed and decision.reason in {"hidden_plugin_skip_auth"}:
        flags.should_return_allowed = True
        return flags

    bot_decision = deps.policy_decision_point.decide_bot(prep.policy_context)
    if bot_decision.allowed:
        hook_recorder.set("auth_bot", "policy")
    elif bot_decision.denied:
        raise_for_policy(bot_decision, deps.policy_skip_message(bot_decision.reason))
    elif bot_decision.deferred:
        raise PermissionExemption(f"auth_bot deferred: {bot_decision.reason}")

    group_decision = deps.policy_decision_point.decide_group(prep.policy_context)
    if group_decision.allowed or group_decision.skipped:
        hook_recorder.set("auth_group", f"policy:{group_decision.reason}")
    elif group_decision.denied:
        raise_for_policy(
            group_decision,
            deps.policy_skip_message(group_decision.reason),
        )
    elif group_decision.deferred:
        raise PermissionExemption(f"auth_group deferred: {group_decision.reason}")

    plugin_decision = deps.policy_decision_point.decide_plugin(prep.policy_context)
    if plugin_decision.allowed or plugin_decision.skipped:
        hook_recorder.set("auth_plugin", f"policy:{plugin_decision.reason}")
    elif plugin_decision.denied:
        raise_for_policy(
            plugin_decision,
            deps.policy_skip_message(plugin_decision.reason),
        )
    else:
        raise PermissionExemption(f"auth_plugin deferred: {plugin_decision.reason}")

    admin_decision = deps.policy_decision_point.decide_admin(prep.policy_context)
    if admin_decision.allowed or admin_decision.skipped:
        hook_recorder.set("auth_admin", f"policy:{admin_decision.reason}")
    elif admin_decision.denied:
        raise_for_policy(
            admin_decision,
            deps.policy_skip_message(admin_decision.reason),
        )
    else:
        raise PermissionExemption(f"auth_admin deferred: {admin_decision.reason}")

    return flags


async def route_gate_stage(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> None:
    if not ctx.module:
        ctx.stop(allowed=True, effect="allow", reason="empty_module")
        return

    side_effect_cache = _side_effect_cache(ctx)
    ctx.side_effect_lock = side_effect_cache.lock_for(ctx.module)
    await ctx.side_effect_lock.acquire()
    ctx.entered_side_effect_lock = True

    auth_result_cache = side_effect_cache.auth_results
    ctx.auth_result_cache = auth_result_cache
    cached_result = auth_result_cache.get(ctx.module)
    if cached_result is not None:
        allowed, reason = cached_result
        if not allowed:
            ctx.decision_effect = "skip"
            ctx.decision_reason = reason or "auth_cached_skip"
            raise SkipPluginException(reason or "auth cached skip")
        ctx.stop(allowed=True, effect="allow", reason="auth_cached_allow")
        return

    if deps.is_hidden_plugin(ctx.matcher):
        ctx.stop(allowed=True, effect="allow", reason="hidden_plugin")
        return
    if (
        ctx.event_cache is not None
        and ctx.event_cache.get("ban_state") is True
        and not ctx.event_context.is_superuser
    ):
        ctx.decision_effect = "skip"
        ctx.decision_reason = "ban_cached"
        raise SkipPluginException("user or group banned (cached)")

    if ctx.route_modules is None:
        ctx.route_modules = await deps.get_route_context(ctx.text, ctx.event_cache)
        set_route_modules(ctx.state, ctx.event_context, ctx.route_modules)
    route_missed = (
        ctx.is_command_matcher
        and ctx.module in deps.route_modules_with_commands
        and ctx.module not in ctx.route_modules
        and not deps.matcher_has_alconna_shortcuts(type(ctx.matcher))
    )
    if route_missed:
        if ctx.event_cache is not None:
            ctx.event_cache["route_miss_after_native_match"] = True
        _recorder(ctx).set("route", "miss")


async def prepare_snapshot_stage(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> None:
    ctx.prep = await deps.prepare_auth_state_with_fallback(
        module=ctx.module,
        context=ctx.event_context,
        bot=ctx.bot,
        event_cache=ctx.event_cache,
        skip_ban=ctx.skip_ban,
        hook_recorder=ctx.hook_recorder,
        state=ctx.state,
        session=ctx.session,
    )
    if ctx.prep is None:
        ctx.stop(allowed=True, effect="allow", reason="prepare_timeout_allow")


async def policy_precheck_stage(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> None:
    try:
        ctx.flags = apply_policy_precheck(ctx, deps)
    except PermissionExemption as exc:
        _recorder(ctx).set("policy_fallback", str(exc))
        if is_db_unhealthy():
            ctx.stop(allowed=True, effect="allow", reason="db_unhealthy_cache_miss")
            return
        ctx.prep = await deps.prepare_auth_state(
            module=ctx.module,
            context=ctx.event_context,
            bot=ctx.bot,
            event_cache=ctx.event_cache,
            skip_ban=ctx.skip_ban,
            hook_recorder=ctx.hook_recorder,
            state=ctx.state,
            session=ctx.session,
            allow_cache_load=True,
        )
        if ctx.prep is None:
            ctx.stop(allowed=True, effect="allow", reason="policy_fallback_timeout")
            return
        try:
            ctx.flags = apply_policy_precheck(ctx, deps)
        except PermissionExemption as fallback_exc:
            _recorder(ctx).set("legacy_pure_auth", str(fallback_exc))
            await deps.legacy_pure_auth_fallback(
                prep=ctx.prep,
                event=ctx.event,
                session=ctx.session,
                text=ctx.text,
            )
            ctx.flags = AuthPolicyFlags()
    flags = _require(ctx.flags, "flags")
    if flags.should_return_allowed:
        ctx.stop(allowed=True, effect="allow", reason="policy_precheck_allow")
        return
    await deps.check_ban_from_snapshot(
        prep=ctx.prep,
        matcher=ctx.matcher,
        event_cache=ctx.event_cache,
        skip_ban=ctx.skip_ban,
        hook_recorder=ctx.hook_recorder,
        session=ctx.session,
    )
    ctx.cost_gold = await deps.resolve_cost_gold(
        prep=ctx.prep,
        hook_recorder=ctx.hook_recorder,
        session=ctx.session,
    )


async def legacy_hook_adapter_stage(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> None:
    prep = _prep(ctx)
    deps.bot_filter(ctx.session, context=prep.permission_context)
    ctx.hooks_time = await deps.run_auth_hooks(
        prep=prep,
        session=ctx.session,
        event_cache=ctx.event_cache,
        lane_context=_lane_context(ctx),
        hook_recorder=_recorder(ctx),
        side_effect_commit=_side_effect_commit(ctx),
    )
    ctx.auth_allowed = True
    ctx.decision_effect = "allow"
    ctx.decision_reason = "auth_passed"


async def side_effect_commit_stage(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> None:
    commit = _side_effect_commit(ctx)
    side_effect_cache = _side_effect_cache(ctx)
    if ctx.ignore_flag:
        await commit.rollback_all("auth_ignored")
        return
    if ctx.cost_gold <= 0:
        if commit.has_pending:
            side_effect_cache.commits[ctx.module] = commit
        return
    gold_start = time.time()
    try:
        reservation = await deps.reserve_gold(
            _entity(ctx).user_id,
            ctx.module,
            ctx.cost_gold,
            ctx.session,
        )
        await commit.reserve_gold(
            reservation,
            amount=ctx.cost_gold,
            metadata={"module": ctx.module},
        )
        _recorder(ctx).set("reserve_gold", f"{time.time() - gold_start:.3f}s")
    except deps.insufficient_gold_error:
        deps.logger.debug(
            f"预扣金币失败，金币不足: {ctx.module}",
            deps.log_command,
            session=ctx.session,
        )
        raise SkipPluginException(f"{ctx.module} 金币不足，已取消执行...") from None
    except TimeoutError:
        deps.logger.error(
            f"预扣金币超时，模块: {ctx.module}",
            deps.log_command,
            session=ctx.session,
        )
        raise
    side_effect_cache.commits[ctx.module] = commit


async def decision_log_stage(
    ctx: AuthPipelineContext,
    deps: AuthPipelineDependencies,
) -> None:
    commit = ctx.side_effect_commit
    has_deferred_commit = commit is not None and commit.has_pending
    if (
        ctx.auth_result_cache is not None
        and ctx.auth_allowed is not None
        and not has_deferred_commit
    ):
        ctx.auth_result_cache[ctx.module] = (
            ctx.auth_allowed,
            None if ctx.auth_allowed else ctx.decision_reason,
        )
    if ctx.entered_side_effect_lock and ctx.side_effect_lock is not None:
        try:
            ctx.side_effect_lock.release()
        except Exception:
            pass
        ctx.entered_side_effect_lock = False


def build_auth_pipeline(deps: AuthPipelineDependencies) -> AuthPipeline:
    return AuthPipeline(
        [
            AuthPipelineStage("route_gate", lambda ctx: route_gate_stage(ctx, deps)),
            AuthPipelineStage(
                "prepare_snapshot",
                lambda ctx: prepare_snapshot_stage(ctx, deps),
            ),
            AuthPipelineStage(
                "policy_precheck",
                lambda ctx: policy_precheck_stage(ctx, deps),
            ),
            AuthPipelineStage(
                "legacy_hook_adapter",
                lambda ctx: legacy_hook_adapter_stage(ctx, deps),
            ),
            AuthPipelineStage(
                "side_effect_commit",
                lambda ctx: side_effect_commit_stage(ctx, deps),
            ),
        ]
    )