from __future__ import annotations import asyncio import json import secrets import time 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_handle: str channel_user_id: str chat_guid: str requested_at: float expires_at: float metadata: dict[str, Any] | None = None class IMessagePairingManager: def __init__(self, storage_dir: str | None = None, account_id: str = "default", config_writes: bool = True): 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" / "imessage") self._account_id = account_id self._config_writes = config_writes def _schedule_save(self) -> None: if not self._config_writes: return try: loop = asyncio.get_running_loop() loop.create_task(self._save_all()) except RuntimeError: pass async def _save_all(self) -> None: await self.save_allowlist() await self.save_pending() def generate_code(self) -> str: return "".join(secrets.choice(PAIRING_CODE_CHARS) for _ in range(PAIRING_CODE_LENGTH)) def request_pairing( self, sender_handle: str, channel_user_id: str, chat_guid: 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"[iMessage/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("[iMessage/Pairing] Max pending requests reached") return None code = self.generate_code() self._pending[code] = PairingRequest( code=code, sender_handle=sender_handle, channel_user_id=channel_user_id, chat_guid=chat_guid, requested_at=time.time(), expires_at=time.time() + PAIRING_VALIDITY_S, metadata=metadata, ) self._rate_limit[channel_user_id] = now logger.info(f"[iMessage/Pairing] New pairing request: code={code}, sender={sender_handle}") self._schedule_save() 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"[iMessage/Pairing] Approved: {req.channel_user_id}") self._schedule_save() return True def reject(self, code: str) -> bool: req = self._pending.pop(code, None) if req is None: return False logger.info(f"[iMessage/Pairing] Rejected: {req.channel_user_id}") self._schedule_save() return True def list_pending(self) -> list[dict[str, Any]]: self._cleanup_expired() return [ { "code": r.code, "sender_handle": r.sender_handle, "channel_user_id": r.channel_user_id, "chat_guid": r.chat_guid, "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 def get_paired_handles(self) -> list[str]: return sorted(self._approved) def is_paired(self, channel_user_id: str) -> bool: return channel_user_id in self._approved def revoke_approval(self, channel_user_id: str) -> bool: if channel_user_id in self._approved: self._approved.discard(channel_user_id) self._schedule_save() return True return False 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) -> None: dir_path = Path(self._storage_dir) / self._account_id dir_path.mkdir(parents=True, exist_ok=True) file_path = dir_path / "paired_handles.json" file_path.write_text( json.dumps(sorted(self._approved), ensure_ascii=False, indent=2), encoding="utf-8", ) async def load_allowlist(self) -> list[str]: file_path = Path(self._storage_dir) / self._account_id / "paired_handles.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"[iMessage/Pairing] Failed to load paired handles: {e}") return [] async def save_pending(self) -> None: dir_path = Path(self._storage_dir) / self._account_id dir_path.mkdir(parents=True, exist_ok=True) pending_data = { code: { "sender_handle": r.sender_handle, "channel_user_id": r.channel_user_id, "chat_guid": r.chat_guid, "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) -> None: file_path = Path(self._storage_dir) / self._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_handle=info.get("sender_handle", ""), channel_user_id=info.get("channel_user_id", ""), chat_guid=info.get("chat_guid", ""), 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"[iMessage/Pairing] Failed to load pending: {e}")