improved with factory pattern (#120)

Layr-Labs · Feb 23, 2024 · d145eb2 · d145eb2
1 parent 597396e
commit d145eb2
Showing 1 changed file with 52 additions and 79 deletions.
diff --git a/cmd/egnkey/generate.go b/cmd/egnkey/generate.go
@@ -63,93 +63,66 @@ It creates the following artifacts based on arguments
 	},
 }
 
-func generate(c *cli.Context) error {
-	keyType := c.String(KeyTypeFlag.Name)
-	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)
-	}
-
-	// TODO: This can be improved further with a factory pattern
-	if keyType == "ecdsa" {
-
-		folder, err := createDir(c, "ecdsa-")
-		if err != nil {
-			return err
-		}
-
-		passwordFile, privateKeyFile, err := createPasswordAndPrivateKeyFiles(folder)
-
-		if err != nil {
-			return err
-		}
-
-		err = generateECDSAKeys(numKeys, folder, passwordFile, privateKeyFile)
-		if err != nil {
-			return err
-		}
-
-	} else if keyType == "bls" {
-
-		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
-		}
+// KeyGenerator defines the interface for generating keys
+type KeyGenerator interface {
+    GenerateKeys(numKeys int, path string, passwordFile, privateKeyFile *os.File) error
+}
 
-		blsFolder, err := createDir(c, "bls-")
-		if err != nil {
-			return err
-		}
+// ECDSAKeyGenerator implements KeyGenerator for ECDSA keys
+type ECDSAKeyGenerator struct{}
 
-		blsPasswordFile, blsPrivateKeyFile, err := createPasswordAndPrivateKeyFiles(blsFolder)
+func (g ECDSAKeyGenerator) GenerateKeys(numKeys int, path string, passwordFile, privateKeyFile *os.File) error {
+    return generateECDSAKeys(numKeys, path, passwordFile, privateKeyFile) // Existing logic from generateECDSAKeys
+}
 
-		if err != nil {
-			return err
-		}
+// BLSKeyGenerator implements KeyGenerator for BLS keys
+type BLSKeyGenerator struct{}
 
-		err = generateBlsKeys(numKeys, blsFolder, blsPasswordFile, blsPrivateKeyFile)
-		if err != nil {
-			return err
-		}
+func (g BLSKeyGenerator) GenerateKeys(numKeys int, path string, passwordFile, privateKeyFile *os.File) error {
+    return generateBlsKeys(numKeys, path, passwordFile, privateKeyFile) // Existing logic from generateBlsKeys
+}
 
-	} else {
-		return cli.Exit("Invalid key type", 1)
-	}
+// NewKeyGenerator is the factory function to create a KeyGenerator based on the key type
+func NewKeyGenerator(keyType string) KeyGenerator {
+    switch keyType {
+    case "ecdsa":
+        return ECDSAKeyGenerator{}
+    case "bls":
+        return BLSKeyGenerator{}
+    default:
+        return nil
+    }
+}
 
-	return nil
+// Modified generate function using the factory pattern
+func generate(c *cli.Context) error {
+    keyType := c.String(KeyTypeFlag.Name)
+    numKeys := c.Int(NumKeysFlag.Name)
+
+    generator := NewKeyGenerator(keyType)
+    if generator == nil {
+        return cli.Exit("Invalid key type", 1)
+    }
+
+    folder, err := createDir(c, keyType+"-")
+    if err != nil {
+        return err
+    }
+
+    passwordFile, privateKeyFile, err := createPasswordAndPrivateKeyFiles(folder)
+    if err != nil {
+        return err
+    }
+
+    err = generator.GenerateKeys(numKeys, folder, passwordFile, privateKeyFile)
+    if err != nil {
+        return err
+    }
+
+    return nil
 }
 
+
 func createDir(c *cli.Context, prefix string) (fileName string, err error) {
 	folder := c.String(OutputDirFlag.Name)
 	if folder == "" {