"""多渠道网关 DM 配对码管理。""" from __future__ import annotations import hmac import secrets from datetime import timedelta from typing import TYPE_CHECKING from sqlalchemy import delete, select from yuxi.channel.constants import InboundRejectionReason from yuxi.channel.plugins.protocol import OutboundMessage from yuxi.storage.postgres.manager import pg_manager from yuxi.storage.postgres.model_channel import ChannelPairingRecord from yuxi.utils.datetime_utils import utc_now_naive from yuxi.utils.logging_config import logger from .policy import SecurityCheckResult from .qr_binding import PairingQRGenerator if TYPE_CHECKING: from .registry import SecurityContext class PairingChecker: """SecurityChecker 包装器:在 DM 白名单拒绝且 mode != deny 时生成配对码。""" name = "pairing" default_priority = 200 def __init__(self, manager: PairingManager | None = None) -> None: self.manager = manager or PairingManager() async def check(self, ctx: SecurityContext) -> SecurityCheckResult | None: peer_id = ctx.inbound.peer_id or ctx.inbound.sender_id or "" pairing_cfg = ctx.config_checker or ctx.config.get("pairing", {}) mode = pairing_cfg.get("mode", "code") try: code = await self.manager.generate_pairing_code( ctx.inbound.channel_type, ctx.inbound.account_id, peer_id, pairing_mode=mode, ) except Exception: logger.exception("Failed to generate pairing code") return SecurityCheckResult( allowed=False, reason=InboundRejectionReason.DM_NOT_ALLOWED, ) if mode in ("qr", "both"): qr_generator = PairingQRGenerator(base_url=pairing_cfg.get("qr_link_base_url")) qr_content = qr_generator.generate_content( code, ctx.inbound.channel_type, ctx.inbound.account_id, peer_id, mode=pairing_cfg.get("qr_content_mode", "plain"), ) build_qr_reply = getattr(ctx.plugin, "build_pairing_qr_reply", None) if build_qr_reply is not None: qr_reply = await build_qr_reply(code, qr_content, ctx.config, ctx.inbound.account_id) else: qr_reply = OutboundMessage(content=qr_content or code, content_type="text") return SecurityCheckResult( allowed=False, reason=InboundRejectionReason.DM_PAIRING_REQUIRED, pairing_code=code, qr_content=qr_content, qr_reply=qr_reply, ) return SecurityCheckResult( allowed=False, reason=InboundRejectionReason.DM_PAIRING_REQUIRED, pairing_code=code, ) class PairingManager: CODE_TTL_MINUTES = 10 CODE_LENGTH = 6 def generate_code_string(self) -> str: """生成一个独立的 6 位数字配对码字符串(不写入数据库)。""" return "".join(secrets.choice("0123456789") for _ in range(self.CODE_LENGTH)) async def generate_pairing_code( self, channel_type: str, account_id: str, peer_id: str, platform_user_id: str | None = None, qr_content: str | None = None, pairing_mode: str = "code", code: str | None = None, ) -> str: code = code or self.generate_code_string() token = secrets.token_urlsafe(32) expires_at = utc_now_naive() + timedelta(minutes=self.CODE_TTL_MINUTES) async with pg_manager.get_async_session_context() as session: await session.execute( delete(ChannelPairingRecord).where( ChannelPairingRecord.channel_type == channel_type, ChannelPairingRecord.account_id == account_id, ChannelPairingRecord.peer_id == peer_id, ChannelPairingRecord.status == "pending", ) ) record = ChannelPairingRecord( channel_type=channel_type, account_id=account_id, peer_id=peer_id, pairing_code=code, pairing_token=token, status="pending", expires_at=expires_at, platform_user_id=platform_user_id, qr_content=qr_content, pairing_mode=pairing_mode, ) session.add(record) await session.commit() return code async def verify_pairing_code( self, channel_type: str, account_id: str, peer_id: str, code: str, ) -> bool: record = await self.verify_pairing_code_record(channel_type, account_id, peer_id, code) return record is not None async def verify_pairing_code_record( self, channel_type: str, account_id: str, peer_id: str, code: str, ) -> ChannelPairingRecord | None: """验证配对码并返回匹配的记录,验证成功时标记为已配对。""" now = utc_now_naive() async with pg_manager.get_async_session_context() as session: result = await session.execute( select(ChannelPairingRecord) .where( ChannelPairingRecord.channel_type == channel_type, ChannelPairingRecord.account_id == account_id, ChannelPairingRecord.peer_id == peer_id, ChannelPairingRecord.status == "pending", ChannelPairingRecord.expires_at > now, ) .order_by(ChannelPairingRecord.created_at.desc()) .limit(1) .with_for_update() ) record = result.scalar_one_or_none() if record is not None and hmac.compare_digest(record.pairing_code, code): record.status = "paired" record.paired_at = now await session.commit() return record return None