import asyncio import hmac import logging import secrets import time from abc import ABC, abstractmethod from dataclasses import dataclass from datetime import UTC, datetime, timedelta from yuxi.repositories.channel_pairing_repo import ChannelPairingRecordRepository PAIRING_CODE_LENGTH = 8 PAIRING_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" PAIRING_CODE_MAX_ATTEMPTS = 500 PAIRING_PENDING_TTL_SECONDS = 600 MAX_PENDING_PER_ACCOUNT = 3 logger = logging.getLogger(__name__) def _generate_code(alphabet: str, length: int) -> str: return "".join(secrets.choice(alphabet) for _ in range(length)) def _generate_token() -> str: return secrets.token_hex(32) @dataclass class PendingPairingRequest: channel_type: str peer_id: str account_id: str code: str created_at: float _record_id: str | None = None class PairingError(Exception): pass class PairingStore(ABC): @abstractmethod async def load_pending(self, channel_type: str, account_id: str) -> list[PendingPairingRequest]: ... @abstractmethod async def upsert_pending(self, req: PendingPairingRequest) -> None: ... @abstractmethod async def remove_pending(self, channel_type: str, peer_id: str, account_id: str) -> None: ... async def expire_pending(self, channel_type: str, peer_id: str, account_id: str) -> None: await self.remove_pending(channel_type, peer_id, account_id) class InMemoryPairingStore(PairingStore): def __init__(self): self._pending: dict[str, dict[str, PendingPairingRequest]] = {} self._lock = asyncio.Lock() @staticmethod def _account_key(channel_type: str, account_id: str) -> str: return f"{channel_type}:{account_id}" async def load_pending(self, channel_type: str, account_id: str) -> list[PendingPairingRequest]: async with self._lock: now = time.monotonic() account_key = self._account_key(channel_type, account_id) bucket = self._pending.get(account_key) if bucket is None: return [] result: list[PendingPairingRequest] = [] expired_peers: list[str] = [] for peer_id, v in bucket.items(): if now - v.created_at >= PAIRING_PENDING_TTL_SECONDS: expired_peers.append(peer_id) continue result.append(v) for peer_id in expired_peers: del bucket[peer_id] if not bucket: del self._pending[account_key] return result async def upsert_pending(self, req: PendingPairingRequest) -> None: account_key = self._account_key(req.channel_type, req.account_id) async with self._lock: bucket = self._pending.setdefault(account_key, {}) bucket[req.peer_id] = req async def remove_pending(self, channel_type: str, peer_id: str, account_id: str) -> None: account_key = self._account_key(channel_type, account_id) async with self._lock: bucket = self._pending.get(account_key) if bucket is None: return bucket.pop(peer_id, None) if not bucket: del self._pending[account_key] class PostgresPairingStore(PairingStore): def __init__(self, repo: ChannelPairingRecordRepository | None = None): self._repo = repo or ChannelPairingRecordRepository() async def load_pending(self, channel_type: str, account_id: str) -> list[PendingPairingRequest]: records = await self._repo.find_pending_by_account(channel_type, account_id) result: list[PendingPairingRequest] = [] for record in records: result.append( PendingPairingRequest( channel_type=record.channel_type, peer_id=record.peer_id, account_id=record.account_id, code=record.pairing_code, created_at=record.created_at.timestamp(), ) ) return result async def upsert_pending(self, req: PendingPairingRequest) -> None: expires_at = datetime.now(UTC).replace(tzinfo=None) + timedelta(seconds=PAIRING_PENDING_TTL_SECONDS) token = _generate_token() record = await self._repo.upsert_pending( channel_type=req.channel_type, account_id=req.account_id, peer_id=req.peer_id, code=req.code, token=token, expires_at=expires_at, ) req._record_id = record.id async def remove_pending(self, channel_type: str, peer_id: str, account_id: str) -> None: await self._repo.remove_pending(channel_type, account_id, peer_id, target_status="paired") async def expire_pending(self, channel_type: str, peer_id: str, account_id: str) -> None: await self._repo.remove_pending(channel_type, account_id, peer_id, target_status="expired") class PairingManager: def __init__(self, store: PairingStore | None = None, allowlist=None): self._store = store or PostgresPairingStore() self._allowlist = allowlist self._verify_attempts: dict[str, list[float]] = {} self._attempts_lock = asyncio.Lock() def _attempts_key(self, channel_type: str, peer_id: str, account_id: str) -> str: return f"{channel_type}:{account_id}:{peer_id}" async def _check_rate_limit(self, key: str) -> bool: async with self._attempts_lock: now = time.monotonic() attempts = self._verify_attempts.get(key, []) attempts = [t for t in attempts if now - t < 300] if len(attempts) >= 10: self._verify_attempts[key] = attempts return False attempts.append(now) self._verify_attempts[key] = attempts self._cleanup_stale_attempts(now) return True def _cleanup_stale_attempts(self, now: float) -> None: stale = [k for k, v in self._verify_attempts.items() if not v or all(now - t >= 300 for t in v)] for k in stale: del self._verify_attempts[k] async def upsert_code(self, channel_type: str, peer_id: str, account_id: str = "default") -> PendingPairingRequest: now = time.monotonic() existing = await self._store.load_pending(channel_type, account_id) existing_req = next((r for r in existing if r.peer_id == peer_id), None) if existing_req is not None: return existing_req if len(existing) >= MAX_PENDING_PER_ACCOUNT: raise RuntimeError(f"Too many pending pairings for channel '{channel_type}' account '{account_id}'") existing_codes = {r.code for r in existing} for _ in range(PAIRING_CODE_MAX_ATTEMPTS): code = _generate_code(PAIRING_CODE_ALPHABET, PAIRING_CODE_LENGTH) if code not in existing_codes: break else: raise RuntimeError("Failed to generate unique pairing code") req = PendingPairingRequest( channel_type=channel_type, peer_id=peer_id, account_id=account_id, code=code, created_at=now, ) await self._store.upsert_pending(req) return req async def verify(self, channel_type: str, peer_id: str, code: str, account_id: str = "default") -> bool: key = self._attempts_key(channel_type, peer_id, account_id) if not await self._check_rate_limit(key): logger.warning("Pairing verify rate limit exceeded for %s", key) return False now = time.monotonic() existing = await self._store.load_pending(channel_type, account_id) for record in existing: if record.peer_id != peer_id: continue if now - record.created_at > PAIRING_PENDING_TTL_SECONDS: await self._store.expire_pending(channel_type, peer_id, account_id) return False if not hmac.compare_digest(record.code, code.upper()): return False try: await self._store.remove_pending(channel_type, peer_id, account_id) except Exception: logger.exception("Failed to remove pending pairing after successful verification") return True return False async def list_pending(self, channel_type: str, account_id: str = "default") -> list[PendingPairingRequest]: return await self._store.load_pending(channel_type, account_id) async def approve(self, channel_type: str, peer_id: str, account_id: str = "default") -> None: if self._allowlist is not None: await self._allowlist.add_entry(channel_type, "dm", peer_id) await self._store.remove_pending(channel_type, peer_id, account_id) async def reject(self, channel_type: str, peer_id: str, account_id: str = "default") -> None: if self._allowlist is not None: await self._allowlist.remove_entry(channel_type, "dm", peer_id) await self._store.expire_pending(channel_type, peer_id, account_id)