ForcePilot/backend/package/yuxi/channel/extensions/nostr/key_utils.py

import re

import bech32
from coincurve import PrivateKey


def validate_private_key(key: str) -> bytes:
    key = key.strip()
    if key.startswith("nsec1"):
        hrp, data = bech32.bech32_decode(key)
        if hrp is None or data is None:
            raise ValueError("Invalid nsec1 format")
        sk_bytes = bytes(bech32.convertbits(data, 5, 8, False))
    elif re.match(r"^[0-9a-fA-F]{64}$", key):
        sk_bytes = bytes.fromhex(key)
    else:
        raise ValueError("Invalid private key format: expected nsec1... or 64-char hex")
    if len(sk_bytes) != 32:
        raise ValueError(f"Private key must be 32 bytes, got {len(sk_bytes)}")
    return sk_bytes


def get_public_key(sk: str | bytes) -> str:
    if isinstance(sk, str):
        sk = validate_private_key(sk)
    pk = PrivateKey(sk)
    pubkey_bytes = pk.public_key.format(compressed=False)
    return pubkey_bytes[1:33].hex()


def normalize_pubkey(pubkey: str) -> str:
    pubkey = pubkey.strip()
    if pubkey.startswith("npub1"):
        hrp, data = bech32.bech32_decode(pubkey)
        if hrp is None or data is None:
            raise ValueError(f"Invalid npub1 format: {pubkey[:20]}...")
        pk_bytes = bytes(bech32.convertbits(data, 5, 8, False))
        return pk_bytes.hex().lower()
    if re.match(r"^[0-9a-fA-F]{64}$", pubkey):
        return pubkey.lower()
    raise ValueError(f"Invalid pubkey format: {pubkey[:20]}...")


def pubkey_to_npub(hex_pk: str) -> str:
    data = bech32.convertbits(bytes.fromhex(hex_pk), 8, 5, True)
    return bech32.bech32_encode("npub", data)


def is_valid_pubkey(pubkey: str) -> bool:
    try:
        normalize_pubkey(pubkey)
        return True
    except (ValueError, IndexError, TypeError):
        return False


def _encode_tlv(items: list[tuple[int, bytes]]) -> bytes:
    result = bytearray()
    for typ, data in items:
        result.append(typ)
        result.extend(len(data).to_bytes(2, "big"))
        result.extend(data)
    return bytes(result)


def _decode_tlv(data: bytes) -> list[tuple[int, bytes]]:
    items = []
    i = 0
    while i < len(data):
        if i + 2 >= len(data):
            break
        typ = data[i]
        length = int.from_bytes(data[i + 1:i + 3], "big")
        value = data[i + 3:i + 3 + length]
        items.append((typ, value))
        i += 3 + length
    return items


def pubkey_to_nprofile(hex_pk: str, relays: list[str] | None = None) -> str:
    items = [(0, bytes.fromhex(hex_pk))]
    if relays:
        for r in relays:
            items.append((1, r.encode("utf-8")))
    data = _encode_tlv(items)
    converted = bech32.convertbits(data, 8, 5, True)
    return bech32.bech32_encode("nprofile", converted)


def event_to_nevent(
    event_id_hex: str, relays: list[str] | None = None,
    author_pk: str | None = None, kind: int | None = None,
) -> str:
    items = [(0, bytes.fromhex(event_id_hex))]
    if relays:
        for r in relays:
            items.append((1, r.encode("utf-8")))
    if author_pk:
        items.append((2, bytes.fromhex(author_pk)))
    if kind is not None:
        items.append((3, kind.to_bytes(4, "big")))
    data = _encode_tlv(items)
    converted = bech32.convertbits(data, 8, 5, True)
    return bech32.bech32_encode("nevent", converted)


def kind_pk_d_to_naddr(kind: int, pubkey: str, d_tag: str = "", relays: list[str] | None = None) -> str:
    items = [(0, d_tag.encode("utf-8"))]
    if relays:
        for r in relays:
            items.append((1, r.encode("utf-8")))
    items.append((2, bytes.fromhex(pubkey)))
    items.append((3, kind.to_bytes(4, "big")))
    data = _encode_tlv(items)
    converted = bech32.convertbits(data, 8, 5, True)
    return bech32.bech32_encode("naddr", converted)


def event_to_note(event_id_hex: str) -> str:
    data = bech32.convertbits(bytes.fromhex(event_id_hex), 8, 5, True)
    return bech32.bech32_encode("note", data)


def decode_nprofile(bech32_str: str) -> tuple[str, list[str]]:
    hrp, data = bech32.bech32_decode(bech32_str)
    if hrp is None or data is None:
        raise ValueError("Invalid nprofile")
    raw = bytes(bech32.convertbits(data, 5, 8, False))
    items = _decode_tlv(raw)
    pubkey = ""
    relays = []
    for typ, value in items:
        if typ == 0:
            pubkey = value.hex()
        elif typ == 1:
            relays.append(value.decode("utf-8"))
    return pubkey, relays


def decode_nevent(bech32_str: str) -> dict:
    hrp, data = bech32.bech32_decode(bech32_str)
    if hrp is None or data is None:
        raise ValueError("Invalid nevent")
    raw = bytes(bech32.convertbits(data, 5, 8, False))
    items = _decode_tlv(raw)
    result: dict = {"event_id": "", "relays": [], "author": None, "kind": None}
    for typ, value in items:
        if typ == 0:
            result["event_id"] = value.hex()
        elif typ == 1:
            result["relays"].append(value.decode("utf-8"))
        elif typ == 2:
            result["author"] = value.hex()
        elif typ == 3:
            result["kind"] = int.from_bytes(value, "big")
    return result