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

from __future__ import annotations

import json
import time

from nostr_sdk import (
    Event,
    EventBuilder,
    Keys,
    Kind,
    PublicKey,
    Tag,
    UnwrappedGift,
    nip04_decrypt,
    nip04_encrypt,
)

from yuxi.utils.logging_config import logger


class NostrCryptoError(Exception):
    pass


def normalize_pubkey(raw: str) -> str:
    if not raw or not raw.strip():
        return ""
    raw = raw.strip()
    if raw.startswith("nostr:"):
        raw = raw[6:]
    if raw.startswith("npub1"):
        try:
            return PublicKey.from_bech32(raw).to_hex()
        except Exception:
            return ""
    if raw.startswith("nprofile1"):
        try:
            from nostr_sdk import Profile

            profile = Profile.from_bech32(raw)
            return profile.public_key().to_hex()
        except Exception:
            return ""
    if len(raw) == 64 and all(c in "0123456789abcdef" for c in raw):
        return raw
    if raw.startswith("nsec1"):
        try:
            keys = Keys.parse(raw)
            return keys.public_key().to_hex()
        except Exception:
            return ""
    return ""


class NostrCrypto:
    def __init__(self, nsec_or_hex: str | None = None):
        if nsec_or_hex:
            try:
                self._keys = Keys.parse(nsec_or_hex)
            except Exception as e:
                msg = f"Nostr 私钥解析失败: {e}。请检查配置中的 private_key (nsec/hex) 是否正确。"
                logger.error(msg)
                raise NostrCryptoError(msg) from e
        else:
            self._keys = Keys.generate()
            logger.warning(
                "未配置 Nostr 私钥，已自动生成临时密钥。npub: %s。生产环境请配置正确的 private_key",
                self._keys.public_key().to_bech32(),
            )
        self._public_key = self._keys.public_key()
        self._secret_key = self._keys.secret_key()

    @property
    def npub(self) -> str:
        return self._public_key.to_bech32()

    def pubkey_hex(self) -> str:
        return self._public_key.to_hex()

    def nsec(self) -> str:
        return self._secret_key.to_bech32()

    def encrypt_nip04(self, plaintext: str, receiver_pubkey_hex: str) -> str:
        receiver_pk = PublicKey.from_hex(receiver_pubkey_hex)
        return nip04_encrypt(self._secret_key, receiver_pk, plaintext)

    def decrypt_nip04(self, ciphertext: str, sender_pubkey_hex: str) -> str:
        sender_pk = PublicKey.from_hex(sender_pubkey_hex)
        return nip04_decrypt(self._secret_key, sender_pk, ciphertext)

    async def encrypt_nip17(self, plaintext: str, receiver_pubkey_hex: str) -> str:
        receiver_pk = PublicKey.from_hex(receiver_pubkey_hex)
        rumor = EventBuilder.private_msg_rumor(receiver_pk, plaintext).build(self._public_key)
        sealed = await EventBuilder.seal(self._keys, receiver_pk, rumor)
        return sealed.sign_with_keys(self._keys).as_json()

    async def decrypt_nip17(self, wrapped_json: str) -> str:
        event = Event.from_json(wrapped_json)
        unwrapped = UnwrappedGift.from_gift_wrap(event)
        if unwrapped is None:
            raise NostrCryptoError("NIP-17 解密失败: Gift Wrap 解封返回空")
        rumor = unwrapped.rumor()
        if rumor is None:
            raise NostrCryptoError("NIP-17 解密失败: Rumor 为空")
        return rumor.content()

    def build_and_sign_event(self, kind: int, content: str, tags: list[list[str]]) -> dict:
        builder = EventBuilder(Kind(kind), content)
        parsed_tags = []
        for tag in tags:
            try:
                parsed_tags.append(Tag.parse(tag))
            except Exception as e:
                raise NostrCryptoError(f"Tag 解析失败: {tag}, 错误: {e}") from e
        if parsed_tags:
            builder = builder.tags(parsed_tags)
        builder = builder.custom_created_at(int(time.time()))
        event = builder.sign_with_keys(self._keys)
        return json.loads(event.as_json())

    def verify_event(self, raw_event: dict) -> bool:
        try:
            event = Event.from_json(json.dumps(raw_event))
            event.verify()
            return True
        except Exception:
            return False

    def verify_event_strict(self, raw_event: dict) -> bool:
        try:
            event = Event.from_json(json.dumps(raw_event))
            event_id = event.id()
            expected_id = raw_event.get("id", "")
            if event_id.to_hex() != expected_id and event_id.to_bech32() != expected_id:
                return False
            event.verify()
            return True
        except Exception:
            return False