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}")