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

"""多渠道网关 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