Add support for generating multiple bls keys for a given shard and node/network

cmd/subcommands/keys.go

@@ -33,6 +33,8 @@ var (
 	passphraseFilePath     string
 	passphrase             string
 	blsFilePath            string
+	blsShardID             uint32
+	blsCount               uint32
 	ppPrompt               = fmt.Sprintf(
 		"prompt for passphrase, otherwise use default passphrase: \"`%s`\"", c.DefaultPassphrase,
 	)
@@ -321,6 +323,24 @@ func keysSub() []*cobra.Command {
 	cmdGenerateBlsKey.Flags().BoolVar(&userProvidesPassphrase, "passphrase", false, ppPrompt)
 	cmdGenerateBlsKey.Flags().StringVar(&passphraseFilePath, "passphrase-file", "", "path to a file containing the passphrase")
 
+	cmdGenerateMultiBlsKeys := &cobra.Command{
+		Use:   "generate-bls-keys",
+		Short: "Generates multiple bls keys for a given shard and then encrypts and saves the private key with a requested passphrase",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			passphrase, err := getPassphraseWithConfirm()
+			if err != nil {
+				return err
+			}
+			return keys.GenMultiBlsKeys(passphrase, blsFilePath, node, blsCount, blsShardID)
+		},
+	}
+	cmdGenerateMultiBlsKeys.Flags().StringVar(&blsFilePath, "bls-file-path", "",
+		"absolute path of where to save encrypted bls private keys")
+	cmdGenerateMultiBlsKeys.Flags().BoolVar(&userProvidesPassphrase, "passphrase", false, ppPrompt)
+	cmdGenerateMultiBlsKeys.Flags().StringVar(&passphraseFilePath, "passphrase-file", "", "path to a file containing the passphrase")
+	cmdGenerateMultiBlsKeys.Flags().Uint32Var(&blsShardID, "shard", 0, "which shard to create bls keys for")
+	cmdGenerateMultiBlsKeys.Flags().Uint32Var(&blsCount, "count", 1, "how many bls keys to generate")
+
 	cmdRecoverBlsKey := &cobra.Command{
 		Use:   "recover-bls-key <ABSOLUTE_PATH_BLS_KEY>",
 		Short: "Recover bls keys from an encrypted bls key file",
@@ -363,7 +383,7 @@ func keysSub() []*cobra.Command {
 	}
 
 	return []*cobra.Command{cmdList, cmdLocation, cmdAdd, cmdRemove, cmdMnemonic, cmdRecoverMnemonic, cmdImportKS, cmdImportPK,
-		cmdExportKS, cmdExportPK, cmdGenerateBlsKey, cmdRecoverBlsKey, cmdSaveBlsKey, GetPublicBlsKey}
+		cmdExportKS, cmdExportPK, cmdGenerateBlsKey, cmdGenerateMultiBlsKeys, cmdRecoverBlsKey, cmdSaveBlsKey, GetPublicBlsKey}
 }
 
 func init() {

pkg/keys/bls.go

@@ -11,12 +11,14 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math/big"
 	"os"
 	"path"
 	"strings"
 
 	ffiBls "github.com/harmony-one/bls/ffi/go/bls"
 	"github.com/harmony-one/go-sdk/pkg/common"
+	"github.com/harmony-one/go-sdk/pkg/sharding"
 	"github.com/harmony-one/harmony/crypto/bls"
 	"github.com/harmony-one/harmony/crypto/hash"
 	"github.com/harmony-one/harmony/shard"
@@ -30,25 +32,40 @@ func GenBlsKeys(passphrase, filePath string) error {
 	publicKeyHex := publicKey.SerializeToHexStr()
 	privateKeyHex := privateKey.SerializeToHexStr()
 
-	if filePath == "" {
-		cwd, _ := os.Getwd()
-		filePath = fmt.Sprintf("%s/%s.key", cwd, publicKeyHex)
-	}
-	if !path.IsAbs(filePath) {
-		return common.ErrNotAbsPath
-	}
-	encryptedPrivateKeyStr, err := encrypt([]byte(privateKeyHex), passphrase)
+	return writeBlsKeyToFile(passphrase, filePath, publicKeyHex, privateKeyHex)
+}
+
+// GenMultiBlsKeys - generate multiple BLS keys for a given shard and node/network
+func GenMultiBlsKeys(passphrase string, filePath string, node string, count uint32, shardID uint32) error {
+	shardingStructure, err := sharding.Structure(node)
 	if err != nil {
 		return err
 	}
-	err = writeToFile(filePath, encryptedPrivateKeyStr)
-	if err != nil {
-		return err
+	shardCount := len(shardingStructure)
+	generatedCount := uint32(0)
+
+	for {
+		if generatedCount < count {
+			privateKey := bls.RandPrivateKey()
+			publicKey := privateKey.GetPublicKey()
+			publicKeyHex := publicKey.SerializeToHexStr()
+			privateKeyHex := privateKey.SerializeToHexStr()
+			shardPubKey := new(shard.BlsPublicKey)
+			if err = shardPubKey.FromLibBLSPublicKey(publicKey); err != nil {
+				return err
+			}
+
+			if blsKeyMatchesShardID(shardPubKey, shardID, shardCount) {
+				if err = writeBlsKeyToFile(passphrase, filePath, publicKeyHex, privateKeyHex); err != nil {
+					return err
+				}
+				generatedCount++
+			}
+		} else {
+			break
+		}
 	}
-	out := fmt.Sprintf(`
-{"public-key" : "%s", "private-key" : "%s", "encrypted-private-key-path" : "%s"}`,
-		publicKeyHex, privateKeyHex, filePath)
-	fmt.Println(common.JSONPrettyFormat(out))
+
 	return nil
 }
 
@@ -191,6 +208,29 @@ func getBlsKey(privateKeyHex string) (*ffiBls.SecretKey, error) {
 	return privateKey, nil
 }
 
+func writeBlsKeyToFile(passphrase, filePath, publicKeyHex, privateKeyHex string) error {
+	if filePath == "" {
+		cwd, _ := os.Getwd()
+		filePath = fmt.Sprintf("%s/%s.key", cwd, publicKeyHex)
+	}
+	if !path.IsAbs(filePath) {
+		return common.ErrNotAbsPath
+	}
+	encryptedPrivateKeyStr, err := encrypt([]byte(privateKeyHex), passphrase)
+	if err != nil {
+		return err
+	}
+	err = writeToFile(filePath, encryptedPrivateKeyStr)
+	if err != nil {
+		return err
+	}
+	out := fmt.Sprintf(`
+{"public-key" : "%s", "private-key" : "%s", "encrypted-private-key-path" : "%s"}`,
+		publicKeyHex, privateKeyHex, filePath)
+	fmt.Println(common.JSONPrettyFormat(out))
+	return nil
+}
+
 func writeToFile(filename string, data string) error {
 	file, err := os.Create(filename)
 	if err != nil {
@@ -257,3 +297,9 @@ func decryptRaw(data []byte, passphrase string) ([]byte, error) {
 	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
 	return plaintext, err
 }
+
+func blsKeyMatchesShardID(pubKey *shard.BlsPublicKey, shardID uint32, shardCount int) bool {
+	bigShardCount := big.NewInt(int64(shardCount))
+	resolvedShardID := int(new(big.Int).Mod(pubKey.Big(), bigShardCount).Int64())
+	return (int(shardID) == resolvedShardID)
+}