ForcePilot/backend/package/yuxi/channels/adapters/wechat/pairing.py

from __future__ import annotations

import json
import secrets
import time
from collections.abc import Callable, Awaitable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from yuxi.utils.logging_config import logger

PAIRING_CODE_CHARS = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"
PAIRING_CODE_LENGTH = 8
PAIRING_VALIDITY_S = 3600
PAIRING_RATE_LIMIT_S = 3600
PAIRING_MAX_PENDING = 3


@dataclass
class PairingRequest:
    code: str
    sender_id: str
    channel_user_id: str
    requested_at: float
    expires_at: float
    metadata: dict[str, Any] | None = None


class WeChatPairingAdapter:
    id_label: str = "WeChat 用户 ID"

    def __init__(self, storage_dir: str | None = None):
        self._pending: dict[str, PairingRequest] = {}
        self._rate_limit: dict[str, float] = {}
        self._approved: set[str] = set()
        self._storage_dir = storage_dir or str(Path.home() / ".yuxi" / "wechat")

    @staticmethod
    def normalize_allow_entry(raw: str) -> str:
        raw = raw.strip()
        if not raw.startswith("wx:"):
            return f"wx:{raw}"
        return raw

    def generate_code(self) -> str:
        return "".join(secrets.choice(PAIRING_CODE_CHARS) for _ in range(PAIRING_CODE_LENGTH))

    def request_pairing(
        self,
        sender_id: str,
        channel_user_id: str,
        metadata: dict[str, Any] | None = None,
    ) -> str | None:
        if self.is_approved(channel_user_id):
            return None

        now = time.monotonic()
        last_request = self._rate_limit.get(channel_user_id, 0)
        if now - last_request < PAIRING_RATE_LIMIT_S:
            logger.info(f"[WeChat/Pairing] Rate limited for {channel_user_id}")
            return None

        if len(self._pending) >= PAIRING_MAX_PENDING:
            self._cleanup_expired()
            if len(self._pending) >= PAIRING_MAX_PENDING:
                logger.warning("[WeChat/Pairing] Max pending requests reached")
                return None

        code = self.generate_code()
        self._pending[code] = PairingRequest(
            code=code,
            sender_id=sender_id,
            channel_user_id=channel_user_id,
            requested_at=time.time(),
            expires_at=time.time() + PAIRING_VALIDITY_S,
            metadata=metadata,
        )
        self._rate_limit[channel_user_id] = now
        logger.info(f"[WeChat/Pairing] New request: code={code}, sender={sender_id}")
        return code

    def approve(self, code: str) -> bool:
        self._cleanup_expired()
        req = self._pending.pop(code, None)
        if req is None:
            return False

        self._approved.add(req.channel_user_id)
        logger.info(f"[WeChat/Pairing] Approved: {req.channel_user_id}")
        return True

    def reject(self, code: str) -> bool:
        req = self._pending.pop(code, None)
        if req is None:
            return False
        logger.info(f"[WeChat/Pairing] Rejected: {req.channel_user_id}")
        return True

    def list_pending(self) -> list[dict[str, Any]]:
        self._cleanup_expired()
        return [
            {
                "code": r.code,
                "sender_id": r.sender_id,
                "channel_user_id": r.channel_user_id,
                "requested_at": r.requested_at,
                "expires_at": r.expires_at,
            }
            for r in self._pending.values()
        ]

    def is_approved(self, channel_user_id: str) -> bool:
        return channel_user_id in self._approved

    async def notify_approval(
        self,
        send_fn: Callable[..., Awaitable[Any]],
        channel_user_id: str,
    ) -> None:
        try:
            await send_fn(channel_user_id, "您的访问请求已获批准，现在可以与我对话了！")
        except Exception as e:
            logger.error(f"[WeChat/Pairing] notify_approval failed: {e}")

    def _cleanup_expired(self) -> None:
        now = time.time()
        expired = [code for code, req in self._pending.items() if req.expires_at < now]
        for code in expired:
            self._pending.pop(code, None)

    async def save_allowlist(self, account_id: str = "default") -> None:
        dir_path = Path(self._storage_dir) / account_id
        dir_path.mkdir(parents=True, exist_ok=True)
        file_path = dir_path / "allowFrom.json"
        file_path.write_text(
            json.dumps(sorted(self._approved), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    async def load_allowlist(self, account_id: str = "default") -> list[str]:
        file_path = Path(self._storage_dir) / account_id / "allowFrom.json"
        if not file_path.exists():
            return []
        try:
            data = json.loads(file_path.read_text(encoding="utf-8"))
            if isinstance(data, list):
                self._approved.update(data)
                return data
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"[WeChat/Pairing] Failed to load allowlist: {e}")
        return []

    async def save_pending(self, account_id: str = "default") -> None:
        dir_path = Path(self._storage_dir) / account_id
        dir_path.mkdir(parents=True, exist_ok=True)
        pending_data = {
            code: {
                "sender_id": r.sender_id,
                "channel_user_id": r.channel_user_id,
                "requested_at": r.requested_at,
                "expires_at": r.expires_at,
            }
            for code, r in self._pending.items()
        }
        file_path = dir_path / "pairing-pending.json"
        file_path.write_text(json.dumps(pending_data, ensure_ascii=False, indent=2), encoding="utf-8")

    async def load_pending(self, account_id: str = "default") -> None:
        file_path = Path(self._storage_dir) / account_id / "pairing-pending.json"
        if not file_path.exists():
            return
        try:
            data = json.loads(file_path.read_text(encoding="utf-8"))
            for code, info in data.items():
                if info.get("expires_at", 0) < time.time():
                    continue
                self._pending[code] = PairingRequest(
                    code=code,
                    sender_id=info.get("sender_id", ""),
                    channel_user_id=info.get("channel_user_id", ""),
                    requested_at=info.get("requested_at", time.time()),
                    expires_at=info.get("expires_at", time.time() + PAIRING_VALIDITY_S),
                )
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"[WeChat/Pairing] Failed to load pending: {e}")