ForcePilot/backend/package/yuxi/channel/security/pairing.py

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)