From 09a626506deb8ce578d8621774b21981219cb029 Mon Sep 17 00:00:00 2001
From: anupsv <anupsv@users.noreply.github.com>
Date: Tue, 1 Oct 2024 16:58:50 -0700
Subject: [PATCH] fixing curve opts (#8)

---
 curve/ops.go         | 13 ++++++++++---
 keystore/keystore.go | 24 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/curve/ops.go b/curve/ops.go
index adc1d34..887a30e 100644
--- a/curve/ops.go
+++ b/curve/ops.go
@@ -2,10 +2,10 @@ package curve
 
 import (
 	"encoding/hex"
+	bls12381Fr "github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
 	"math/big"
 
 	bls12381 "github.com/consensys/gnark-crypto/ecc/bls12-381"
-	bls12381Fr "github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
 	"github.com/consensys/gnark-crypto/ecc/bn254"
 	bn254Fr "github.com/consensys/gnark-crypto/ecc/bn254/fr"
 )
@@ -18,8 +18,9 @@ const (
 )
 
 type Ops struct {
-	GenerateG2PubKey func(secret []byte) string
-	GenerateG1PubKey func(secret []byte) string
+	GenerateG2PubKey            func(secret []byte) string
+	GenerateG1PubKey            func(secret []byte) string
+	GeneratePrivateKeyFromBytes func(secret []byte) [32]byte
 }
 
 var OpsMap = map[Curve]Ops{
@@ -46,6 +47,9 @@ var OpsMap = map[Curve]Ops{
 			publicKeyBytes := pubKey.ScalarMultiplication(&g1Gen, &frBigInt).Bytes()
 			return hex.EncodeToString(publicKeyBytes[:])
 		},
+		GeneratePrivateKeyFromBytes: func(secret []byte) [32]byte {
+			return new(bls12381Fr.Element).SetBytes(secret).Bytes()
+		},
 	},
 	BN254: {
 		GenerateG2PubKey: func(secret []byte) string {
@@ -70,5 +74,8 @@ var OpsMap = map[Curve]Ops{
 			publicKeyBytes := pubKey.ScalarMultiplication(&g1Gen, &frBigInt).Bytes()
 			return hex.EncodeToString(publicKeyBytes[:])
 		},
+		GeneratePrivateKeyFromBytes: func(secret []byte) [32]byte {
+			return new(bn254Fr.Element).SetBytes(secret).Bytes()
+		},
 	},
 }
diff --git a/keystore/keystore.go b/keystore/keystore.go
index 5359073..2db69db 100644
--- a/keystore/keystore.go
+++ b/keystore/keystore.go
@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	curveOps "github.com/Layr-Labs/bn254-keystore-go/curve"
 	"os"
 	"path/filepath"
 	"unicode"
@@ -510,6 +511,29 @@ func NewKeyPair(
 	}, nil
 }
 
+// NewKeyPairWithCurve generates a new key pair using the provided password and mnemonic language and the curve.
+func NewKeyPairWithCurve(
+	password string,
+	language mnemonic.Language,
+	curve curve.Curve,
+) (*KeyPair, error) {
+
+	ops, exists := curveOps.OpsMap[curve]
+	if !exists {
+		return nil, fmt.Errorf("curve '%s' not supported", curve)
+	}
+
+	keyPair, err := NewKeyPair(password, language)
+
+	if err != nil {
+		return nil, err
+	}
+
+	privateKeyFrBytes := ops.GeneratePrivateKeyFromBytes(keyPair.PrivateKey)
+	keyPair.PrivateKey = privateKeyFrBytes[:]
+	return keyPair, nil
+}
+
 func NewKeyPairFromMnemonic(
 	pkMnemonic string,
 	password string,