feat: add NIP-04 encrypt/decrypt methods on Key

2 files changed, 198 insertions, 0 deletions

diff --git a/nip04.go b/nip04.go
new file mode 100644
index 0000000..1f1c245
--- /dev/null
+++ b/nip04.go
@@ -0,0 +1,149 @@
+package nostr
+
+import (
+        "crypto/aes"
+        "crypto/cipher"
+        "crypto/rand"
+        "crypto/sha256"
+        "encoding/base64"
+        "encoding/hex"
+        "fmt"
+        "strings"
+
+        "github.com/btcsuite/btcd/btcec/v2"
+)
+
+// NIP04Encrypt encrypts a plaintext message for a recipient using NIP-04
+// (ECDH shared secret + AES-256-CBC). The sender's private key and
+// recipient's public key (64-char hex) are required.
+// Returns the ciphertext in NIP-04 format: "<base64>?iv=<base64>"
+func (k *Key) NIP04Encrypt(recipientPubHex, plaintext string) (string, error) {
+        if k.priv == nil {
+                return "", fmt.Errorf("cannot encrypt: public-only key")
+        }
+
+        shared, err := nip04SharedSecret(k.priv, recipientPubHex)
+        if err != nil {
+                return "", err
+        }
+
+        iv := make([]byte, 16)
+        if _, err := rand.Read(iv); err != nil {
+                return "", fmt.Errorf("generating iv: %w", err)
+        }
+
+        block, err := aes.NewCipher(shared)
+        if err != nil {
+                return "", fmt.Errorf("aes cipher: %w", err)
+        }
+
+        padded := pkcs7Pad([]byte(plaintext), aes.BlockSize)
+        ct := make([]byte, len(padded))
+        cipher.NewCBCEncrypter(block, iv).CryptBlocks(ct, padded)
+
+        return base64.StdEncoding.EncodeToString(ct) + "?iv=" + base64.StdEncoding.EncodeToString(iv), nil
+}
+
+// NIP04Decrypt decrypts a NIP-04 ciphertext using the receiver's private key
+// and sender's public key (64-char hex).
+func (k *Key) NIP04Decrypt(senderPubHex, ciphertext string) (string, error) {
+        if k.priv == nil {
+                return "", fmt.Errorf("cannot decrypt: public-only key")
+        }
+
+        parts := strings.SplitN(ciphertext, "?iv=", 2)
+        if len(parts) != 2 {
+                return "", fmt.Errorf("invalid NIP-04 ciphertext format")
+        }
+
+        ct, err := base64.StdEncoding.DecodeString(parts[0])
+        if err != nil {
+                return "", fmt.Errorf("decode ciphertext: %w", err)
+        }
+
+        iv, err := base64.StdEncoding.DecodeString(parts[1])
+        if err != nil {
+                return "", fmt.Errorf("decode iv: %w", err)
+        }
+
+        shared, err := nip04SharedSecret(k.priv, senderPubHex)
+        if err != nil {
+                return "", err
+        }
+
+        block, err := aes.NewCipher(shared)
+        if err != nil {
+                return "", fmt.Errorf("aes cipher: %w", err)
+        }
+
+        if len(ct)%aes.BlockSize != 0 {
+                return "", fmt.Errorf("ciphertext is not a multiple of block size")
+        }
+
+        plain := make([]byte, len(ct))
+        cipher.NewCBCDecrypter(block, iv).CryptBlocks(plain, ct)
+
+        unpadded, err := pkcs7Unpad(plain)
+        if err != nil {
+                return "", fmt.Errorf("unpad: %w", err)
+        }
+
+        return string(unpadded), nil
+}
+
+// nip04SharedSecret computes the ECDH shared secret per NIP-04:
+// SHA256(privKey * recipientPubKey).x
+func nip04SharedSecret(priv *btcec.PrivateKey, recipientPubHex string) ([]byte, error) {
+        pubBytes, err := hex.DecodeString(recipientPubHex)
+        if err != nil {
+                return nil, fmt.Errorf("decode pubkey: %w", err)
+        }
+
+        // NIP-04 pubkeys are x-only (32 bytes, Schnorr). Prefix with 0x02 for
+        // compressed SEC1 parsing.
+        if len(pubBytes) == 32 {
+                pubBytes = append([]byte{0x02}, pubBytes...)
+        }
+
+        pub, err := btcec.ParsePubKey(pubBytes)
+        if err != nil {
+                return nil, fmt.Errorf("parse pubkey: %w", err)
+        }
+
+        // ECDH: scalar multiply
+        var point btcec.JacobianPoint
+        pub.AsJacobian(&point)
+        priv.Key.SetByteSlice(priv.Serialize())
+        btcec.ScalarMultNonConst(&priv.Key, &point, &point)
+        point.ToAffine()
+
+        // Shared secret is SHA256 of the x-coordinate
+        xBytes := point.X.Bytes()
+        shared := sha256.Sum256(xBytes[:])
+
+        return shared[:], nil
+}
+
+func pkcs7Pad(data []byte, blockSize int) []byte {
+        pad := blockSize - len(data)%blockSize
+        padded := make([]byte, len(data)+pad)
+        copy(padded, data)
+        for i := len(data); i < len(padded); i++ {
+                padded[i] = byte(pad)
+        }
+        return padded
+}
+
+func pkcs7Unpad(data []byte) ([]byte, error) {
+        if len(data) == 0 {
+                return nil, fmt.Errorf("empty data")
+        }
+        pad := int(data[len(data)-1])
+        if pad == 0 || pad > aes.BlockSize {
+                return nil, fmt.Errorf("invalid padding size: %d", pad)
+        }
+        if len(data) < pad {
+                return nil, fmt.Errorf("data shorter than padding")
+        }
+        return data[:len(data)-pad], nil
+}
diff --git a/nip04_test.go b/nip04_test.go
new file mode 100644
index 0000000..ff06151
--- /dev/null
+++ b/nip04_test.go
@@ -0,0 +1,49 @@
+package nostr
+
+import (
+        "testing"
+)
+
+func TestNIP04RoundTrip(t *testing.T) {
+        alice, err := GenerateKey()
+        if err != nil {
+                t.Fatal(err)
+        }
+        bob, err := GenerateKey()
+        if err != nil {
+                t.Fatal(err)
+        }
+
+        plaintext := "hello, private world!"
+
+        ct, err := alice.NIP04Encrypt(bob.Public(), plaintext)
+        if err != nil {
+                t.Fatalf("encrypt: %v", err)
+        }
+
+        got, err := bob.NIP04Decrypt(alice.Public(), ct)
+        if err != nil {
+                t.Fatalf("decrypt: %v", err)
+        }
+
+        if got != plaintext {
+                t.Fatalf("expected %q, got %q", plaintext, got)
+        }
+}
+
+func TestNIP04DifferentKeys(t *testing.T) {
+        alice, _ := GenerateKey()
+        bob, _ := GenerateKey()
+        eve, _ := GenerateKey()
+
+        ct, err := alice.NIP04Encrypt(bob.Public(), "secret")
+        if err != nil {
+                t.Fatal(err)
+        }
+
+        // Eve should not be able to decrypt
+        _, err = eve.NIP04Decrypt(alice.Public(), ct)
+        if err == nil {
+                t.Fatal("expected error decrypting with wrong key, got nil")
+        }
+}