"""多渠道网关安全策略编排。""" 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)