Add bech32 encoding (npub/nsec for Nostr)

2 files changed, 405 insertions, 0 deletions

diff --git a/bech32.go b/bech32.go
new file mode 100644
index 0000000..835156d
--- /dev/null
+++ b/bech32.go
@@ -0,0 +1,212 @@
+package secp256k1
+
+import (
+        "fmt"
+        "math/big"
+        "strings"
+)
+
+// Bech32 alphabet (no 1, b, i, o to avoid confusion)
+const bech32Alphabet = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
+
+// bech32Polymod computes the Bech32 checksum
+func bech32Polymod(values []int) int {
+        gen := []int{0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3}
+        chk := 1
+        for _, v := range values {
+                top := chk >> 25
+                chk = (chk&0x1ffffff)<<5 ^ v
+                for i := 0; i < 5; i++ {
+                        if (top>>i)&1 == 1 {
+                                chk ^= gen[i]
+                        }
+                }
+        }
+        return chk
+}
+
+// bech32HRPExpand expands the human-readable part for checksum
+func bech32HRPExpand(hrp string) []int {
+        result := make([]int, len(hrp)*2+1)
+        for i, c := range hrp {
+                result[i] = int(c) >> 5
+        }
+        result[len(hrp)] = 0
+        for i, c := range hrp {
+                result[len(hrp)+1+i] = int(c) & 31
+        }
+        return result
+}
+
+// bech32CreateChecksum creates a 6-character checksum
+func bech32CreateChecksum(hrp string, data []int) []int {
+        values := append(bech32HRPExpand(hrp), data...)
+        values = append(values, []int{0, 0, 0, 0, 0, 0}...)
+        polymod := bech32Polymod(values) ^ 1
+        checksum := make([]int, 6)
+        for i := 0; i < 6; i++ {
+                checksum[i] = (polymod >> (5 * (5 - i))) & 31
+        }
+        return checksum
+}
+
+// bech32VerifyChecksum verifies a bech32 checksum
+func bech32VerifyChecksum(hrp string, data []int) bool {
+        return bech32Polymod(append(bech32HRPExpand(hrp), data...)) == 1
+}
+
+// convertBits converts between bit groups
+func convertBits(data []byte, fromBits, toBits int, pad bool) ([]int, error) {
+        acc := 0
+        bits := 0
+        result := []int{}
+        maxv := (1 << toBits) - 1
+
+        for _, value := range data {
+                acc = (acc << fromBits) | int(value)
+                bits += fromBits
+                for bits >= toBits {
+                        bits -= toBits
+                        result = append(result, (acc>>bits)&maxv)
+                }
+        }
+
+        if pad {
+                if bits > 0 {
+                        result = append(result, (acc<<(toBits-bits))&maxv)
+                }
+        } else if bits >= fromBits || ((acc<<(toBits-bits))&maxv) != 0 {
+                return nil, fmt.Errorf("invalid padding")
+        }
+
+        return result, nil
+}
+
+// Bech32Encode encodes data with a human-readable prefix
+func Bech32Encode(hrp string, data []byte) (string, error) {
+        // Convert 8-bit bytes to 5-bit groups
+        values, err := convertBits(data, 8, 5, true)
+        if err != nil {
+                return "", err
+        }
+
+        // Add checksum
+        checksum := bech32CreateChecksum(hrp, values)
+        combined := append(values, checksum...)
+
+        // Build result string
+        var result strings.Builder
+        result.WriteString(hrp)
+        result.WriteByte('1')
+        for _, v := range combined {
+                result.WriteByte(bech32Alphabet[v])
+        }
+
+        return result.String(), nil
+}
+
+// Bech32Decode decodes a bech32 string
+func Bech32Decode(s string) (string, []byte, error) {
+        // Find separator
+        pos := strings.LastIndex(s, "1")
+        if pos < 1 || pos+7 > len(s) {
+                return "", nil, fmt.Errorf("invalid bech32: separator position")
+        }
+
+        // Split HRP and data
+        hrp := strings.ToLower(s[:pos])
+        dataStr := strings.ToLower(s[pos+1:])
+
+        // Decode data characters
+        data := make([]int, len(dataStr))
+        for i, c := range dataStr {
+                idx := strings.IndexRune(bech32Alphabet, c)
+                if idx < 0 {
+                        return "", nil, fmt.Errorf("invalid bech32: character %c", c)
+                }
+                data[i] = idx
+        }
+
+        // Verify checksum
+        if !bech32VerifyChecksum(hrp, data) {
+                return "", nil, fmt.Errorf("invalid bech32: checksum")
+        }
+
+        // Remove checksum and convert back to 8-bit
+        values := data[:len(data)-6]
+        bytes, err := convertBitsToBytes(values)
+        if err != nil {
+                return "", nil, err
+        }
+
+        return hrp, bytes, nil
+}
+
+// convertBitsToBytes converts 5-bit groups back to bytes
+func convertBitsToBytes(data []int) ([]byte, error) {
+        acc := 0
+        bits := 0
+        result := []byte{}
+
+        for _, value := range data {
+                acc = (acc << 5) | value
+                bits += 5
+                for bits >= 8 {
+                        bits -= 8
+                        result = append(result, byte((acc>>bits)&0xff))
+                }
+        }
+
+        return result, nil
+}
+
+// === Nostr-specific helpers ===
+
+// Nsec returns the private key as a Nostr nsec string
+func (priv *PrivateKey) Nsec() string {
+        encoded, _ := Bech32Encode("nsec", priv.Bytes())
+        return encoded
+}
+
+// Npub returns the public key as a Nostr npub string (x-only)
+func (pub *PublicKey) Npub() string {
+        encoded, _ := Bech32Encode("npub", pub.XOnlyBytes())
+        return encoded
+}
+
+// PrivateKeyFromNsec parses an nsec string
+func PrivateKeyFromNsec(nsec string) (*PrivateKey, error) {
+        hrp, data, err := Bech32Decode(nsec)
+        if err != nil {
+                return nil, fmt.Errorf("invalid nsec: %w", err)
+        }
+        if hrp != "nsec" {
+                return nil, fmt.Errorf("invalid nsec: wrong prefix %q", hrp)
+        }
+        if len(data) != 32 {
+                return nil, fmt.Errorf("invalid nsec: wrong length %d", len(data))
+        }
+        return NewPrivateKeyFromBytes(data)
+}
+
+// PublicKeyFromNpub parses an npub string
+func PublicKeyFromNpub(npub string) (*PublicKey, error) {
+        hrp, data, err := Bech32Decode(npub)
+        if err != nil {
+                return nil, fmt.Errorf("invalid npub: %w", err)
+        }
+        if hrp != "npub" {
+                return nil, fmt.Errorf("invalid npub: wrong prefix %q", hrp)
+        }
+        if len(data) != 32 {
+                return nil, fmt.Errorf("invalid npub: wrong length %d", len(data))
+        }
+
+        // Lift x-coordinate to full point
+        point, err := liftX(new(big.Int).SetBytes(data))
+        if err != nil {
+                return nil, fmt.Errorf("invalid npub: %w", err)
+        }
+
+        return &PublicKey{Point: point}, nil
+}
diff --git a/bech32_test.go b/bech32_test.go
new file mode 100644
index 0000000..276d998
--- /dev/null
+++ b/bech32_test.go
@@ -0,0 +1,193 @@
+package secp256k1
+
+import (
+        "strings"
+        "testing"
+)
+
+func TestBech32EncodeBasic(t *testing.T) {
+        data := []byte{0x00, 0x01, 0x02}
+        encoded, err := Bech32Encode("test", data)
+        if err != nil {
+                t.Fatalf("encoding failed: %v", err)
+        }
+
+        // Should start with hrp + "1"
+        if !strings.HasPrefix(encoded, "test1") {
+                t.Errorf("encoded should start with 'test1', got %s", encoded)
+        }
+}
+
+func TestBech32RoundTrip(t *testing.T) {
+        data := []byte{0xde, 0xad, 0xbe, 0xef}
+        encoded, err := Bech32Encode("test", data)
+        if err != nil {
+                t.Fatalf("encoding failed: %v", err)
+        }
+
+        hrp, decoded, err := Bech32Decode(encoded)
+        if err != nil {
+                t.Fatalf("decoding failed: %v", err)
+        }
+
+        if hrp != "test" {
+                t.Errorf("hrp mismatch: got %s, want test", hrp)
+        }
+
+        if len(decoded) != len(data) {
+                t.Fatalf("length mismatch: got %d, want %d", len(decoded), len(data))
+        }
+
+        for i := range data {
+                if decoded[i] != data[i] {
+                        t.Errorf("byte %d mismatch: got %x, want %x", i, decoded[i], data[i])
+                }
+        }
+}
+
+func TestBech32DecodeInvalidChecksum(t *testing.T) {
+        // Valid encoding, then corrupt it
+        data := []byte{0x01, 0x02, 0x03}
+        encoded, _ := Bech32Encode("test", data)
+
+        // Corrupt last character
+        corrupted := encoded[:len(encoded)-1] + "q"
+
+        _, _, err := Bech32Decode(corrupted)
+        if err == nil {
+                t.Error("should reject invalid checksum")
+        }
+}
+
+func TestBech32DecodeInvalidCharacter(t *testing.T) {
+        _, _, err := Bech32Decode("test1invalid!")
+        if err == nil {
+                t.Error("should reject invalid character")
+        }
+}
+
+func TestNsecEncode(t *testing.T) {
+        priv, _ := NewPrivateKeyFromHex("0000000000000000000000000000000000000000000000000000000000000001")
+        nsec := priv.Nsec()
+
+        if !strings.HasPrefix(nsec, "nsec1") {
+                t.Errorf("nsec should start with 'nsec1', got %s", nsec)
+        }
+}
+
+func TestNpubEncode(t *testing.T) {
+        priv, _ := NewPrivateKeyFromHex("0000000000000000000000000000000000000000000000000000000000000001")
+        pub := priv.PublicKey()
+        npub := pub.Npub()
+
+        if !strings.HasPrefix(npub, "npub1") {
+                t.Errorf("npub should start with 'npub1', got %s", npub)
+        }
+}
+
+func TestNsecRoundTrip(t *testing.T) {
+        priv1, _ := GeneratePrivateKey()
+        nsec := priv1.Nsec()
+
+        priv2, err := PrivateKeyFromNsec(nsec)
+        if err != nil {
+                t.Fatalf("failed to parse nsec: %v", err)
+        }
+
+        if priv1.D.Cmp(priv2.D) != 0 {
+                t.Error("private key should survive nsec round-trip")
+        }
+}
+
+func TestNpubRoundTrip(t *testing.T) {
+        priv, _ := GeneratePrivateKey()
+        pub1 := priv.PublicKey()
+        npub := pub1.Npub()
+
+        pub2, err := PublicKeyFromNpub(npub)
+        if err != nil {
+                t.Fatalf("failed to parse npub: %v", err)
+        }
+
+        // X coordinates should match (y might differ in sign)
+        if pub1.Point.x.value.Cmp(pub2.Point.x.value) != 0 {
+                t.Error("public key x should survive npub round-trip")
+        }
+}
+
+func TestPrivateKeyFromNsecInvalid(t *testing.T) {
+        // Wrong prefix
+        _, err := PrivateKeyFromNsec("npub1qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq3xjaeh")
+        if err == nil {
+                t.Error("should reject npub as nsec")
+        }
+}
+
+func TestPublicKeyFromNpubInvalid(t *testing.T) {
+        // Wrong prefix
+        _, err := PublicKeyFromNpub("nsec1qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq0dcpx3")
+        if err == nil {
+                t.Error("should reject nsec as npub")
+        }
+}
+
+// Test with known private key
+func TestKnownNostrKeyPair(t *testing.T) {
+        // Private key = 1, should give G as public key
+        priv, _ := NewPrivateKeyFromHex("0000000000000000000000000000000000000000000000000000000000000001")
+
+        nsec := priv.Nsec()
+        // Verify it starts with nsec and round-trips
+        if nsec[:5] != "nsec1" {
+                t.Errorf("nsec should start with nsec1, got %s", nsec)
+        }
+
+        // Parse it back
+        priv2, err := PrivateKeyFromNsec(nsec)
+        if err != nil {
+                t.Fatalf("failed to parse nsec: %v", err)
+        }
+        if priv.D.Cmp(priv2.D) != 0 {
+                t.Error("nsec round-trip failed")
+        }
+
+        // Public key should be G
+        pub := priv.PublicKey()
+        npub := pub.Npub()
+
+        if npub[:5] != "npub1" {
+                t.Errorf("npub should start with npub1, got %s", npub)
+        }
+
+        // Parse it back and verify x matches G
+        pub2, err := PublicKeyFromNpub(npub)
+        if err != nil {
+                t.Fatalf("failed to parse npub: %v", err)
+        }
+        if pub2.Point.x.value.Cmp(Gx) != 0 {
+                t.Error("npub should decode to G.x")
+        }
+}
+
+func TestSignAndVerifyWithNostrKeys(t *testing.T) {
+        // Create keys
+        priv, _ := GeneratePrivateKey()
+        nsec := priv.Nsec()
+        npub := priv.PublicKey().Npub()
+
+        // Parse them back
+        priv2, _ := PrivateKeyFromNsec(nsec)
+        pub2, _ := PublicKeyFromNpub(npub)
+
+        // Sign with parsed private key
+        message := []byte("hello nostr")
+        sig, err := Sign(priv2, message)
+        if err != nil {
+                t.Fatalf("signing failed: %v", err)
+        }
+
+        // Verify with parsed public key
+        if !Verify(pub2, message, sig) {
+                t.Error("signature should verify with parsed keys")
+        }
+}