[EgnKey] Generate ECDSA and BLS keys with both-key-type (#61)

* Minor fixes and support generating both as 3rd option * add minor comment * fix pr comment * Fix PR Comment * Fix PR Comments * Fix PR comment --------- Co-authored-by: Siddharth More <Siddhi More>
Layr-Labs · Oct 31, 2023 · 28e8426 · 28e8426
1 parent c396a15
commit 28e8426
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 28 deletions.
diff --git a/cmd/egnkey/README.md b/cmd/egnkey/README.md
@@ -26,4 +26,9 @@ To create in a random folder
 To create  in specific folder
 ```bash
    egnkey generate --key-type ecdsa --num-keys <num_key> --output-dir <path_to_folder>
+```
+
+To create `ECDSA` and `BLS` keys in a random folder
+```bash
+   egnkey generate --key-type both --num-keys <num_key>
 ```
diff --git a/cmd/egnkey/generate.go b/cmd/egnkey/generate.go
@@ -65,59 +65,129 @@ It creates the following artifacts based on arguments
 
 func generate(c *cli.Context) error {
 	keyType := c.String(KeyTypeFlag.Name)
-	if keyType != "ecdsa" && keyType != "bls" {
+	if keyType != "ecdsa" && keyType != "bls" && keyType != "both" {
 		return cli.Exit("Invalid key type", 1)
 	}
 	numKeys := c.Int(NumKeysFlag.Name)
 	if numKeys < 1 {
 		return cli.Exit("Invalid number of keys", 1)
 	}
 
-	folder := c.String(OutputDirFlag.Name)
-	if folder == "" {
-		id, err := uuid.NewUUID()
+	// TODO: This can be improved further with a factory pattern
+	if keyType == "ecdsa" {
+
+		folder, err := createDir(c, "ecdsa-")
 		if err != nil {
-			return cli.Exit("Failed to generate UUID", 1)
+			return err
 		}
-		folder = id.String()
-	}
-
-	_, err := os.Stat(folder)
-	if !os.IsNotExist(err) {
-		return cli.Exit("Folder already exists, choose a different folder or delete the existing folder", 1)
-	}
 
-	err = os.MkdirAll(folder+"/"+DefaultKeyFolder, 0755)
-	if err != nil {
-		return err
-	}
+		passwordFile, privateKeyFile, err := createPasswordAndPrivateKeyFiles(folder)
 
-	passwordFile, err := os.Create(filepath.Clean(folder + "/" + PasswordFile))
-	if err != nil {
-		return err
-	}
-	privateKeyFile, err := os.Create(filepath.Clean(folder + "/" + PrivateKeyHexFile))
-	if err != nil {
-		return err
-	}
+		if err != nil {
+			return err
+		}
 
-	if keyType == "ecdsa" {
-		err := generateECDSAKeys(numKeys, folder, passwordFile, privateKeyFile)
+		err = generateECDSAKeys(numKeys, folder, passwordFile, privateKeyFile)
 		if err != nil {
 			return err
 		}
+
 	} else if keyType == "bls" {
-		err := generateBlsKeys(numKeys, folder, passwordFile, privateKeyFile)
+
+		folder, err := createDir(c, "bls-")
 		if err != nil {
 			return err
 		}
+
+		passwordFile, privateKeyFile, err := createPasswordAndPrivateKeyFiles(folder)
+
+		if err != nil {
+			return err
+		}
+
+		err = generateBlsKeys(numKeys, folder, passwordFile, privateKeyFile)
+		if err != nil {
+			return err
+		}
+	} else if keyType == "both" {
+
+		ecdsaFolder, err := createDir(c, "ecdsa-")
+		if err != nil {
+			return err
+		}
+
+		ecdsaPasswordFile, ecdsaPrivateKeyFile, err := createPasswordAndPrivateKeyFiles(ecdsaFolder)
+
+		if err != nil {
+			return err
+		}
+
+		err = generateECDSAKeys(numKeys, ecdsaFolder, ecdsaPasswordFile, ecdsaPrivateKeyFile)
+		if err != nil {
+			return err
+		}
+
+		blsFolder, err := createDir(c, "bls-")
+		if err != nil {
+			return err
+		}
+
+		blsPasswordFile, blsPrivateKeyFile, err := createPasswordAndPrivateKeyFiles(blsFolder)
+
+		if err != nil {
+			return err
+		}
+
+		err = generateBlsKeys(numKeys, blsFolder, blsPasswordFile, blsPrivateKeyFile)
+		if err != nil {
+			return err
+		}
+
 	} else {
 		return cli.Exit("Invalid key type", 1)
 	}
 
 	return nil
 }
 
+func createDir(c *cli.Context, prefix string) (fileName string, err error) {
+	folder := c.String(OutputDirFlag.Name)
+	if folder == "" {
+		id, err := uuid.NewUUID()
+		if err != nil {
+			return "", err
+		}
+		folder = prefix + id.String()
+	}
+
+	_, err = os.Stat(folder)
+	if !os.IsNotExist(err) {
+		return "", err
+	}
+
+	// Clean the path
+	cleanFilePath := filepath.Clean(folder + "/" + DefaultKeyFolder)
+
+	err = os.MkdirAll(cleanFilePath, 0755)
+	if err != nil {
+		return "", err
+	}
+	return folder, nil
+}
+
+func createPasswordAndPrivateKeyFiles(folder string) (passwordFile, privateKeyFile *os.File, err error) {
+
+	passwordFile, err = os.Create(filepath.Clean(folder + "/" + PasswordFile))
+	if err != nil {
+		return nil, nil, err
+	}
+	privateKeyFile, err = os.Create(filepath.Clean(folder + "/" + PrivateKeyHexFile))
+	if err != nil {
+		return nil, nil, err
+	}
+	return passwordFile, privateKeyFile, nil
+}
+
 func generateBlsKeys(numKeys int, path string, passwordFile, privateKeyFile *os.File) error {
 	for i := 0; i < numKeys; i++ {
 		key, err := bls.GenRandomBlsKeys()
@@ -146,13 +216,28 @@ func generateBlsKeys(numKeys int, path string, passwordFile, privateKeyFile *os.
 		if err != nil {
 			return err
 		}
+
+		if (i+1)%50 == 0 {
+			fmt.Printf("Generated %d keys\n", i+1)
+		}
 	}
 	return nil
 }
 
 func generateECDSAKeys(numKeys int, path string, passwordFile, privateKeyFile *os.File) error {
 	for i := 0; i < numKeys; i++ {
 		key, err := crypto.GenerateKey()
+		privateKeyHex := hex.EncodeToString(key.D.Bytes())
+
+		// Check if the length of privateKeyHex is 32 bytes (64 characters)
+		lenPrivateKey := len(privateKeyHex)
+		if lenPrivateKey != 64 {
+			fmt.Printf("Private key Ignore: %s %d\n", privateKeyHex, lenPrivateKey)
+			// Reset count
+			i--
+			continue
+		}
+
 		if err != nil {
 			return err
 		}
@@ -162,7 +247,6 @@ func generateECDSAKeys(numKeys int, path string, passwordFile, privateKeyFile *o
 			return err
 		}
 
-		privateKeyHex := hex.EncodeToString(key.D.Bytes())
 		fileName := fmt.Sprintf("%d.ecdsa.key.json", i+1)
 		err = ecdsa.WriteKey(filepath.Clean(path+"/"+DefaultKeyFolder+"/"+fileName), key, password)
 		if err != nil {
@@ -178,6 +262,10 @@ func generateECDSAKeys(numKeys int, path string, passwordFile, privateKeyFile *o
 		if err != nil {
 			return err
 		}
+
+		if (i+1)%50 == 0 {
+			fmt.Printf("Generated %d keys\n", i+1)
+		}
 	}
 	return nil
 }