from __future__ import annotations import asyncio import json import secrets import time from dataclasses import dataclass, field from pathlib import Path from typing import Any from yuxi.utils.logging_config import logger PAIRING_CODE_CHARS = "ABCDEFGHJKMNPQRSTUVWXYZ23456789" @dataclass class PairingRequest: code: str sender_id: str channel_user_id: str display_name: str requested_at: float expires_at: float metadata: dict[str, Any] = field(default_factory=dict) @dataclass class PairingConfig: code_length: int = 8 code_ttl_s: float = 3600.0 rate_limit_s: float = 3600.0 max_pending: int = 3 storage_path: str = "" store_pending_file: str = "pairing_pending.json" store_allowlist_file: str = "pairing_allowlist.json" class BasePairingManager: def __init__(self, config: PairingConfig): self._config = config self._pending: dict[str, PairingRequest] = {} self._sender_cooldown: dict[str, float] = {} self._approved: set[str] = set() self._lock = asyncio.Lock() storage_dir = Path(config.storage_path) storage_dir.mkdir(parents=True, exist_ok=True) def generate_code(self) -> str: return "".join(secrets.choice(PAIRING_CODE_CHARS) for _ in range(self._config.code_length)) async def request_pairing( self, sender_id: str, channel_user_id: str, display_name: str = "", metadata: dict[str, Any] | None = None, ) -> str | None: now = time.time() async with self._lock: last_time = self._sender_cooldown.get(sender_id, 0) if now - last_time < self._config.rate_limit_s: logger.debug(f"[Pairing] Rate limited for sender {sender_id}") return None if channel_user_id in self._approved: return None self._cleanup_expired(now) if len(self._pending) >= self._config.max_pending: logger.warning(f"[Pairing] Max pending ({self._config.max_pending}) reached") return None code = self.generate_code() self._pending[code] = PairingRequest( code=code, sender_id=sender_id, channel_user_id=channel_user_id, display_name=display_name, requested_at=now, expires_at=now + self._config.code_ttl_s, metadata=metadata or {}, ) self._sender_cooldown[sender_id] = now await self._save_pending() logger.info(f"[Pairing] New request: {display_name} ({channel_user_id}) — code {code}") return code async def approve(self, code: str) -> bool: async with self._lock: request = self._pending.pop(code, None) if request is None: return False self._approved.add(request.channel_user_id) await self._save_pending() await self._save_allowlist() await self._notify_approval(request.sender_id, request.channel_user_id) logger.info(f"[Pairing] Approved: {request.display_name} ({request.channel_user_id})") return True async def reject(self, code: str) -> bool: async with self._lock: removed = self._pending.pop(code, None) if removed is not None: await self._save_pending() return removed is not None def is_approved(self, channel_user_id: str) -> bool: return channel_user_id in self._approved def list_pending(self) -> list[dict[str, Any]]: now = time.time() return [ { "code": req.code, "display_name": req.display_name, "channel_user_id": req.channel_user_id, "requested_at": req.requested_at, "expires_in_s": max(0, int(req.expires_at - now)), } for req in self._pending.values() if req.expires_at > now ] def get_allowlist(self) -> list[str]: return sorted(self._approved) async def add_to_allowlist(self, channel_user_id: str) -> None: async with self._lock: self._approved.add(channel_user_id) await self._save_allowlist() async def remove_from_allowlist(self, channel_user_id: str) -> bool: async with self._lock: if channel_user_id in self._approved: self._approved.discard(channel_user_id) await self._save_allowlist() return True return False async def load_state(self) -> None: async with self._lock: await self._load_pending() await self._load_allowlist() logger.info(f"[Pairing] Loaded state: {len(self._pending)} pending, {len(self._approved)} approved") async def _save_pending(self) -> None: data = {} now = time.time() for code, req in self._pending.items(): if req.expires_at > now: data[code] = { "sender_id": req.sender_id, "channel_user_id": req.channel_user_id, "display_name": req.display_name, "requested_at": req.requested_at, "expires_at": req.expires_at, "metadata": req.metadata, } path = Path(self._config.storage_path) / self._config.store_pending_file path.write_text(json.dumps(data, ensure_ascii=False, indent=2)) async def _load_pending(self) -> None: path = Path(self._config.storage_path) / self._config.store_pending_file if not path.exists(): return try: data = json.loads(path.read_text()) now = time.time() loaded = 0 for code, item in data.items(): if item["expires_at"] > now: self._pending[code] = PairingRequest( code=code, sender_id=item["sender_id"], channel_user_id=item["channel_user_id"], display_name=item.get("display_name", ""), requested_at=item["requested_at"], expires_at=item["expires_at"], metadata=item.get("metadata", {}), ) loaded += 1 logger.debug(f"[Pairing] Loaded {loaded} pending requests (expired filtered)") except (json.JSONDecodeError, KeyError) as e: logger.warning(f"[Pairing] Failed to load pending: {e}") async def _save_allowlist(self) -> None: path = Path(self._config.storage_path) / self._config.store_allowlist_file path.write_text(json.dumps(sorted(self._approved), ensure_ascii=False, indent=2)) async def _load_allowlist(self) -> None: path = Path(self._config.storage_path) / self._config.store_allowlist_file if not path.exists(): return try: loaded = set(json.loads(path.read_text())) self._approved = loaded logger.debug(f"[Pairing] Loaded {len(loaded)} approved users") except json.JSONDecodeError as e: logger.warning(f"[Pairing] Failed to load allowlist: {e}") def _cleanup_expired(self, now: float) -> None: expired = [code for code, req in self._pending.items() if req.expires_at < now] for code in expired: del self._pending[code] async def _notify_approval(self, sender_id: str, channel_user_id: str) -> None: pass