Kez/python/kez/keys.py

"""Signing keys: nostr (secp256k1 Schnorr) and ed25519."""

from __future__ import annotations

import hashlib
import os

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.serialization import Encoding, PrivateFormat, NoEncryption, PublicFormat

from . import bech32, schnorr
from .identity import Identity
from .jcs import canonical_bytes

NOSTR_SCHNORR_ALG = "nostr-secp256k1-schnorr-sha256-jcs"
ED25519_SHA512_ALG = "ed25519-sha512-jcs"


class NostrSecret:
    """A nostr secp256k1 secret key, addressed by its npub."""

    __slots__ = ("_sk", "_pub")

    def __init__(self, secret_key: bytes) -> None:
        if len(secret_key) != 32:
            raise ValueError("nostr secret key must be 32 bytes")
        self._sk = secret_key
        self._pub = schnorr.pubkey_gen(secret_key)

    @classmethod
    def generate(cls) -> "NostrSecret":
        while True:
            sk = os.urandom(32)
            i = int.from_bytes(sk, "big")
            if 1 <= i < schnorr.n:
                return cls(sk)

    @classmethod
    def from_nsec(cls, nsec: str) -> "NostrSecret":
        raw = bech32.decode("nsec", nsec.strip())
        if len(raw) != 32:
            raise ValueError("invalid nsec: expected 32-byte key")
        return cls(raw)

    def nsec(self) -> str:
        return bech32.encode("nsec", self._sk)

    def npub(self) -> str:
        return bech32.encode("npub", self._pub)

    def identity(self) -> Identity:
        return Identity.parse(f"nostr:{self.npub()}")

    def alg(self) -> str:
        return NOSTR_SCHNORR_ALG

    def sign_payload(self, payload) -> str:
        digest = hashlib.sha256(canonical_bytes(payload)).digest()
        return schnorr.sign(digest, self._sk).hex()


class Ed25519Secret:
    """An ed25519 secret seed, addressed by its hex public key."""

    __slots__ = ("_seed", "_key", "_pub")

    def __init__(self, seed: bytes) -> None:
        if len(seed) != 32:
            raise ValueError("ed25519 seed must be 32 bytes")
        self._seed = seed
        self._key = Ed25519PrivateKey.from_private_bytes(seed)
        self._pub = self._key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)

    @classmethod
    def generate(cls) -> "Ed25519Secret":
        key = Ed25519PrivateKey.generate()
        seed = key.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
        return cls(seed)

    @classmethod
    def from_seed_hex(cls, seed_hex: str) -> "Ed25519Secret":
        seed = bytes.fromhex(seed_hex.strip())
        if len(seed) != 32:
            raise ValueError("invalid ed25519 seed: expected 32-byte (64 hex char) seed")
        return cls(seed)

    @classmethod
    def from_mnemonic(cls, phrase: str) -> "Ed25519Secret":
        # Lazy import: mnemonic.py imports Ed25519Secret at module top.
        from .mnemonic import seed_from_mnemonic

        return cls(seed_from_mnemonic(phrase))

    @classmethod
    def generate_with_mnemonic(cls, words: int = 24) -> tuple["Ed25519Secret", str]:
        from .mnemonic import generate_ed25519_with_mnemonic

        return generate_ed25519_with_mnemonic(words)

    def seed_hex(self) -> str:
        return self._seed.hex()

    def pubkey_hex(self) -> str:
        return self._pub.hex()

    def identity(self) -> Identity:
        return Identity.parse(f"ed25519:{self.pubkey_hex()}")

    def alg(self) -> str:
        return ED25519_SHA512_ALG

    def sign_payload(self, payload) -> str:
        return self._key.sign(canonical_bytes(payload)).hex()


def verify_signature(payload, alg: str, key: Identity, sig_hex: str) -> bool:
    """Verify a signature block against an arbitrary JCS-canonicalizable payload."""
    try:
        sig = bytes.fromhex(sig_hex)
    except ValueError:
        return False
    if len(sig) != 64:
        return False

    if alg == NOSTR_SCHNORR_ALG:
        if key.scheme != "nostr":
            return False
        pubkey = bech32.decode("npub", key.value)
        digest = hashlib.sha256(canonical_bytes(payload)).digest()
        return schnorr.verify(digest, pubkey, sig)

    if alg == ED25519_SHA512_ALG:
        if key.scheme != "ed25519":
            return False
        try:
            pub = Ed25519PublicKey.from_public_bytes(bytes.fromhex(key.value))
            pub.verify(sig, canonical_bytes(payload))
            return True
        except Exception:
            return False

    return False


def signer_from_flags(
    nsec: str | None,
    ed25519_seed: str | None,
    mnemonic: str | None = None,
):
    provided = [v for v in (nsec, ed25519_seed, mnemonic) if v]
    if len(provided) > 1:
        raise ValueError("--nsec, --ed25519-seed, and --mnemonic are mutually exclusive")
    if nsec:
        return NostrSecret.from_nsec(nsec)
    if ed25519_seed:
        return Ed25519Secret.from_seed_hex(ed25519_seed)
    if mnemonic:
        return Ed25519Secret.from_mnemonic(mnemonic)
    raise ValueError("missing key: pass --nsec, --ed25519-seed, or --mnemonic")