Adding encryption of a private key to the key generation.

flare-foundation · Jun 10, 2024 · 9a1a3c8 · 9a1a3c8
1 parent f1a7ac2
commit 9a1a3c8
Show file tree

Hide file tree

Showing 2 changed files with 129 additions and 13 deletions.
diff --git a/go-client/README.md b/go-client/README.md
@@ -152,11 +152,11 @@ go run keygen/keygen.go --key 0x1512de600a10a0aac01580dbfc080965b89ed2329a7b2bf5
 
 where the key value needs to be replaced by the generated private key and the address value needs
 to be replaced by the actual address that will be used to sign the updates.
-Alternatively, one can save and read the key from a file and save the signature with:
+Alternatively, one can save and read the (encrypted) key from a file and save the signature with:
 
 ```bash
-go run keygen/keygen.go --key_out keys.out
-go run keygen/keygen.go --key_file keys.out --address 0xd4e934C2749CA8C1618659D02E7B28B074bf4df7 --sig_out sig.out
+go run keygen/keygen.go --key_out keys.out --pass secret_password
+go run keygen/keygen.go --key_file keys.out --pass secret_password --address 0xd4e934C2749CA8C1618659D02E7B28B074bf4df7 --sig_out sig.out
 ```
 
 where the address value needs to be replaced by the actual address that will be used to sign

diff --git a/go-client/keygen/keygen.go b/go-client/keygen/keygen.go
@@ -1,6 +1,9 @@
 package main
 
 import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/hex"
 	"encoding/json"
@@ -15,12 +18,14 @@ import (
 
 	"github.com/consensys/gnark-crypto/ecc/bn254"
 	"github.com/consensys/gnark-crypto/ecc/bn254/fp"
+	"golang.org/x/crypto/scrypt"
 )
 
 var InFlag = flag.String("key", "", "Private key")
 var AddressFlag = flag.String("address", "", "Value of the address that needs to be signed.")
 var InFileFlag = flag.String("key_file", "", "File to load private and public key")
 var KeyOutFlag = flag.String("key_out", "", "File to save a freshly generated private and public key")
+var PassFlag = flag.String("pass", "", "Password for encrypting/decrypting private key")
 var SigOutFlag = flag.String("sig_out", "", "File to save a signature")
 
 type keyStrings struct {
@@ -29,6 +34,12 @@ type keyStrings struct {
 	PrivateKey string
 }
 
+type keyEncrypted struct {
+	PublicKeyX          string
+	PublicKeyY          string
+	EncryptedPrivateKey []byte
+}
+
 type sigStrings struct {
 	RX string
 	RY string
@@ -50,15 +61,30 @@ func main() {
 		if err != nil {
 			log.Fatal(err)
 		}
-		keyStrings := keyStrings{PrivateKey: "0x" + keys.Sk.Text(16), PublicKeyX: "0x" + keys.Pk.X.Text(16), PublicKeyY: "0x" + keys.Pk.Y.Text(16)}
-		keyBytes, err := json.Marshal(keyStrings)
-		if err != nil {
-			log.Fatal(err)
-		}
 
 		if *KeyOutFlag == "" {
+			keysString := keyStrings{PrivateKey: "0x" + keys.Sk.Text(16), PublicKeyX: "0x" + keys.Pk.X.Text(16), PublicKeyY: "0x" + keys.Pk.Y.Text(16)}
+			keyBytes, err := json.Marshal(keysString)
+			if err != nil {
+				log.Fatal(err)
+			}
 			logger.Info("Key generated: " + string(keyBytes))
 		} else {
+			if *PassFlag == "" {
+				log.Fatal("password should be specified to encrypt the private key")
+			}
+
+			encryptedPrivateKey, err := Encrypt([]byte(*PassFlag), keys.Sk.Bytes())
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			keyEncrypted := keyEncrypted{EncryptedPrivateKey: encryptedPrivateKey, PublicKeyX: "0x" + keys.Pk.X.Text(16), PublicKeyY: "0x" + keys.Pk.Y.Text(16)}
+			keyBytes, err := json.Marshal(keyEncrypted)
+			if err != nil {
+				log.Fatal(err)
+			}
+
 			f, err := os.Create(*KeyOutFlag)
 			if err != nil {
 				log.Fatal(err)
@@ -77,28 +103,37 @@ func main() {
 
 	} else {
 		if *InFileFlag != "" {
+			if *PassFlag == "" {
+				log.Fatal("password should be specified to decrypt the private key")
+			}
+
 			keyBytes, err := os.ReadFile(*InFileFlag)
 			if err != nil {
 				log.Fatal(err)
 			}
-			var keyStrings keyStrings
-			err = json.Unmarshal(keyBytes, &keyStrings)
+			var keyEncrypted keyEncrypted
+			err = json.Unmarshal(keyBytes, &keyEncrypted)
 			if err != nil {
 				log.Fatal(err)
 			}
 
-			keys, err = sortition.KeyFromString(keyStrings.PrivateKey)
+			privateKeyBytes, err := Decrypt([]byte(*PassFlag), keyEncrypted.EncryptedPrivateKey)
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			keys, err = sortition.KeyFromString("0x" + new(big.Int).SetBytes(privateKeyBytes).Text(16))
 			if err != nil {
 				log.Fatal(err)
 			}
 
 			pkCheck := &bn254.G1Affine{}
-			pkXCheck, check := new(big.Int).SetString(keyStrings.PublicKeyX, 0)
+			pkXCheck, check := new(big.Int).SetString(keyEncrypted.PublicKeyX, 0)
 			if !check {
 				log.Fatal(fmt.Errorf("failed to read the key"))
 			}
 			pkCheck.X = *new(fp.Element).SetBigInt(pkXCheck)
-			pkYCheck, check := new(big.Int).SetString(keyStrings.PublicKeyY, 0)
+			pkYCheck, check := new(big.Int).SetString(keyEncrypted.PublicKeyY, 0)
 			if !check {
 				log.Fatal(fmt.Errorf("failed to read the key"))
 			}
@@ -159,6 +194,87 @@ func main() {
 			}
 			logger.Info("Saved the signature in file " + *SigOutFlag)
 		}
+	} else {
+		// in case that key was read and decrypted from a file, but not used for the signature, just print it
+		if *InFileFlag != "" {
+			keysString := keyStrings{PrivateKey: "0x" + keys.Sk.Text(16), PublicKeyX: "0x" + keys.Pk.X.Text(16), PublicKeyY: "0x" + keys.Pk.Y.Text(16)}
+			keyBytes, err := json.Marshal(keysString)
+			if err != nil {
+				log.Fatal(err)
+			}
+			logger.Info("Key: " + string(keyBytes))
+		}
+	}
+}
+
+func Encrypt(password, data []byte) ([]byte, error) {
+	key, salt, err := DeriveKey(password, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	blockCipher, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(blockCipher)
+	if err != nil {
+		return nil, err
+	}
+
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err = rand.Read(nonce); err != nil {
+		return nil, err
+	}
+
+	ciphertext := gcm.Seal(nonce, nonce, data, nil)
+
+	ciphertext = append(ciphertext, salt...)
+
+	return ciphertext, nil
+}
+
+func Decrypt(password, data []byte) ([]byte, error) {
+	salt, data := data[len(data)-32:], data[:len(data)-32]
+
+	key, _, err := DeriveKey(password, salt)
+	if err != nil {
+		return nil, err
+	}
+
+	blockCipher, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(blockCipher)
+	if err != nil {
+		return nil, err
+	}
+
+	nonce, ciphertext := data[:gcm.NonceSize()], data[gcm.NonceSize():]
+
+	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return plaintext, nil
+}
+
+func DeriveKey(password, salt []byte) ([]byte, []byte, error) {
+	if salt == nil {
+		salt = make([]byte, 32)
+		if _, err := rand.Read(salt); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	key, err := scrypt.Key(password, salt, 1048576, 8, 1, 32)
+	if err != nil {
+		return nil, nil, err
 	}
 
+	return key, salt, nil
 }