ForcePilot/backend/package/yuxi/channel/security/policy.py

"""多渠道网关安全策略编排。"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

from yuxi.channel.constants import channel_rate_limit_key
from yuxi.channel.plugins.protocol import OutboundMessage
from yuxi.channel.security.models import DmPolicy

if TYPE_CHECKING:
    from yuxi.channel.message.models import InboundMessage
    from yuxi.channel.plugins.protocol import ChannelPlugin
    from yuxi.storage.postgres.model_channel import ChannelPairingRecord

    from .identity import IdentityLinkResolver
    from .registry import SecurityCheckerRegistry, SecurityContext


@dataclass
class SecurityCheckResult:
    allowed: bool
    reason: str | None = None
    pairing_code: str | None = None
    qr_content: str | None = None
    qr_reply: OutboundMessage | None = None


class SecurityPolicy:
    def __init__(
        self,
        checker_registry: SecurityCheckerRegistry,
        identity_link_resolver: IdentityLinkResolver | None = None,
    ):
        self.checker_registry = checker_registry
        self.identity = identity_link_resolver

    async def check(
        self,
        config: dict,
        plugin: ChannelPlugin,
        inbound: InboundMessage,
    ) -> SecurityCheckResult:
        chat_type = inbound.chat_type or "private"

        resolved_sender_id: str | None = None
        if self.identity is not None and inbound.sender_id:
            resolved_sender_id = self.identity.resolve(inbound.channel_type, inbound.sender_id)

        chain = self.checker_registry.resolve_chain(config)
        if config.get("security_checkers"):
            return await self._run_custom_chain(config, plugin, inbound, chat_type, resolved_sender_id, chain)
        return await self._run_default_chain(config, plugin, inbound, chat_type, resolved_sender_id, chain)

    async def _run_default_chain(
        self,
        config: dict,
        plugin: ChannelPlugin,
        inbound: InboundMessage,
        chat_type: str,
        resolved_sender_id: str | None,
        chain: list,
    ) -> SecurityCheckResult:
        allowlist = self.checker_registry.get("allowlist")
        if allowlist is None:
            return await self._run_custom_chain(config, plugin, inbound, chat_type, resolved_sender_id, chain)

        ctx = self._make_context(config, plugin, inbound, chat_type, resolved_sender_id, allowlist.name)
        result = await allowlist.check(ctx)
        if result is not None and not result.allowed:
            if chat_type == "private":
                dm_policy = self._resolve_dm_policy(config, plugin, inbound)
                if dm_policy.mode != "deny":
                    pairing = self.checker_registry.get("pairing")
                    if pairing is not None:
                        pairing_ctx = self._make_context(
                            config, plugin, inbound, chat_type, resolved_sender_id, pairing.name
                        )
                        return await pairing.check(pairing_ctx)
            return result

        for checker in chain:
            if checker.name in ("allowlist", "pairing"):
                continue
            ctx = self._make_context(config, plugin, inbound, chat_type, resolved_sender_id, checker.name)
            result = await checker.check(ctx)
            if result is not None and not result.allowed:
                return result

        return SecurityCheckResult(allowed=True)

    async def _run_custom_chain(
        self,
        config: dict,
        plugin: ChannelPlugin,
        inbound: InboundMessage,
        chat_type: str,
        resolved_sender_id: str | None,
        chain: list,
    ) -> SecurityCheckResult:
        for checker in chain:
            ctx = self._make_context(config, plugin, inbound, chat_type, resolved_sender_id, checker.name)
            result = await checker.check(ctx)
            if result is not None and not result.allowed:
                return result
        return SecurityCheckResult(allowed=True)

    def _make_context(
        self,
        config: dict,
        plugin: ChannelPlugin,
        inbound: InboundMessage,
        chat_type: str,
        resolved_sender_id: str | None,
        checker_name: str,
    ) -> SecurityContext:
        from .registry import SecurityContext

        config_checker: dict = {}
        for entry in config.get("security_checkers", []):
            if isinstance(entry, dict) and entry.get("name") == checker_name:
                config_checker = entry.get("config") or {}
                break

        return SecurityContext(
            config=config,
            config_checker=config_checker,
            plugin=plugin,
            inbound=inbound,
            chat_type=chat_type,
            resolved_sender_id=resolved_sender_id,
        )

    def _resolve_dm_policy(self, config: dict, plugin: ChannelPlugin, inbound: InboundMessage) -> DmPolicy:
        resolve = getattr(plugin, "resolve_dm_policy", None)
        if resolve is not None:
            policy = resolve(config, inbound.account_id)
            if policy is not None:
                return policy
        return DmPolicy()

    def check_bot_loop(self, actor_id: str, is_bot: bool) -> bool:
        """供扫码事件使用的 bot_loop 检查快捷方法。"""
        checker = self.checker_registry.get("bot_loop")
        if checker is None:
            raise RuntimeError("bot_loop checker not registered")
        return checker.is_allowed(actor_id, is_bot)

    async def check_rate_limit(
        self,
        config: dict,
        plugin: ChannelPlugin,
        inbound: InboundMessage,
        actor_id: str,
    ) -> bool:
        """供扫码事件使用的 rate_limit 检查快捷方法。"""
        checker = self.checker_registry.get("rate_limit")
        if checker is None:
            raise RuntimeError("rate_limit checker not registered")

        resolve = getattr(plugin, "resolve_rate_limit_policy", None)
        rate_policy = resolve(config, inbound.account_id) if resolve is not None else None
        max_requests = (
            rate_policy.max_requests_per_minute
            if rate_policy
            else config.get("rate_limit", {}).get("max_requests_per_minute", 60)
        )
        rate_key = channel_rate_limit_key(inbound.channel_type, inbound.account_id or "", actor_id)
        return await checker.is_allowed(rate_key, max_requests, window_seconds=60)

    async def verify_pairing_code_record(
        self,
        channel_type: str,
        account_id: str,
        peer_id: str,
        code: str,
    ) -> ChannelPairingRecord | None:
        """供扫码事件使用的配对码验证快捷方法。"""
        checker = self.checker_registry.get("pairing")
        if checker is None:
            raise RuntimeError("pairing checker not registered")
        return await checker.manager.verify_pairing_code_record(channel_type, account_id, peer_id, code)