ForcePilot/backend/package/yuxi/channels/adapters/nostr/pair.py

from __future__ import annotations

import secrets
import time
from collections.abc import Awaitable, Callable

from yuxi.channels.adapters.nostr.crypto import NostrCrypto
from yuxi.channels.adapters.nostr.relay_manager import RelayManager
from yuxi.utils.logging_config import logger


class NostrPairingManager:
    PAIRING_CHALLENGE_PREFIX = "🔐 PairingChallenge::"

    def __init__(self, crypto: NostrCrypto, relay_manager: RelayManager):
        self._crypto = crypto
        self._relay_manager = relay_manager
        self._pending: dict[str, dict] = {}
        self._on_pairing_request: list[Callable[[str, str], Awaitable[None]]] = []
        self._on_pairing_approved: list[Callable[[str, str], Awaitable[None]]] = []
        self._on_pairing_denied: list[Callable[[str, str], Awaitable[None]]] = []

    def on_pairing_request(self, handler: Callable[[str, str], Awaitable[None]]) -> None:
        """注册配对请求回调 (challenge_id, pubkey)"""
        self._on_pairing_request.append(handler)

    def on_pairing_approved(self, handler: Callable[[str, str], Awaitable[None]]) -> None:
        """注册配对批准回调 (challenge_id, pubkey)"""
        self._on_pairing_approved.append(handler)

    def on_pairing_denied(self, handler: Callable[[str, str], Awaitable[None]]) -> None:
        """注册配对拒绝回调 (challenge_id, pubkey)"""
        self._on_pairing_denied.append(handler)

    async def _notify_pairing_request(self, challenge_id: str, pubkey: str) -> None:
        for handler in self._on_pairing_request:
            try:
                await handler(challenge_id, pubkey)
            except Exception:
                logger.debug("on_pairing_request handler error", exc_info=True)

    async def _notify_pairing_approved(self, challenge_id: str, pubkey: str) -> None:
        for handler in self._on_pairing_approved:
            try:
                await handler(challenge_id, pubkey)
            except Exception:
                logger.debug("on_pairing_approved handler error", exc_info=True)

    async def _notify_pairing_denied(self, challenge_id: str, pubkey: str) -> None:
        for handler in self._on_pairing_denied:
            try:
                await handler(challenge_id, pubkey)
            except Exception:
                logger.debug("on_pairing_denied handler error", exc_info=True)

    def issue_challenge(self, target_pubkey: str) -> str:
        token = secrets.token_hex(16)
        challenge_id = secrets.token_hex(4)
        full_token = f"{self.PAIRING_CHALLENGE_PREFIX}{challenge_id}::{token}"
        self._pending[challenge_id] = {
            "pubkey": target_pubkey,
            "token": token,
            "status": "pending",
            "created_at": time.monotonic(),
        }
        return full_token

    def verify_response(self, content: str) -> tuple[bool, str]:
        if not content.startswith(self.PAIRING_CHALLENGE_PREFIX):
            return False, "not a pairing response"
        try:
            body = content[len(self.PAIRING_CHALLENGE_PREFIX) :]
            parts = body.split("::", 1)
            challenge_id = parts[0]
            token = parts[1] if len(parts) > 1 else ""
        except (ValueError, IndexError):
            return False, "invalid pairing format"

        pending = self._pending.get(challenge_id)
        if not pending:
            return False, f"unknown challenge id: {challenge_id}"
        if pending["token"] != token:
            return False, "token mismatch"
        if pending["status"] != "pending":
            return False, f"challenge already {pending['status']}"

        pending["status"] = "approved"
        return True, challenge_id

    async def approve_challenge(self, challenge_id: str) -> tuple[bool, str]:
        """通过 UI 回调批准配对请求"""
        pending = self._pending.get(challenge_id)
        if not pending:
            return False, f"unknown challenge id: {challenge_id}"
        if pending["status"] != "pending":
            return False, f"challenge already {pending['status']}"

        pending["status"] = "approved"
        pubkey = pending["pubkey"]
        await self._notify_pairing_approved(challenge_id, pubkey)
        return True, challenge_id

    def deny_challenge(self, challenge_id: str) -> bool:
        pending = self._pending.get(challenge_id)
        if not pending:
            return False
        pending["status"] = "denied"
        return True

    async def deny_challenge_async(self, challenge_id: str) -> bool:
        """通过 UI 回调拒绝配对请求"""
        pending = self._pending.get(challenge_id)
        if not pending:
            return False
        pending["status"] = "denied"
        await self._notify_pairing_denied(challenge_id, pending["pubkey"])
        return True

    def list_pending(self) -> list[dict]:
        return [
            {
                "id": cid,
                "pubkey": p["pubkey"],
                "status": p["status"],
                "created_at": p.get("created_at", 0),
            }
            for cid, p in self._pending.items()
            if p["status"] == "pending"
        ]

    def list_all(self) -> list[dict]:
        return [
            {
                "id": cid,
                "pubkey": p["pubkey"],
                "status": p["status"],
                "created_at": p.get("created_at", 0),
            }
            for cid, p in self._pending.items()
        ]

    def get_pending(self, challenge_id: str) -> dict | None:
        pending = self._pending.get(challenge_id)
        if not pending:
            return None
        return {
            "id": challenge_id,
            "pubkey": pending["pubkey"],
            "status": pending["status"],
            "created_at": pending.get("created_at", 0),
        }

    def cleanup_expired(self, ttl_sec: int = 300) -> int:
        now = time.monotonic()
        removed = 0
        for cid in list(self._pending.keys()):
            entry = self._pending[cid]
            if entry["status"] != "pending":
                self._pending.pop(cid, None)
                removed += 1
            elif now - entry.get("created_at", 0) > ttl_sec:
                entry["status"] = "expired"
                self._pending.pop(cid, None)
                removed += 1
        return removed

    async def send_pairing_challenge(self, target_pubkey: str, chat_type: str = "direct") -> str | None:
        challenge = self.issue_challenge(target_pubkey)
        event = self._crypto.build_and_sign_event(kind=4, content=challenge, tags=[["p", target_pubkey]])
        success_count = await self._relay_manager.broadcast(event)
        if success_count > 0:
            return challenge
        return None