ForcePilot/backend/package/yuxi/channels/adapters/imessage/pairing.py

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