optimize the gen key util

Author: ysmood

lib/secure/secure_test.go

@@ -234,35 +234,26 @@ func TestBelongs(t *testing.T) {
 func TestGenerateKeyFile(t *testing.T) {
 	g := got.T(t)
 
-	g.MkdirAll(0, "tmp")
-
-	p := "tmp/id_ed25519"
-
-	g.E(secure.GenerateKeyFile(false, p, "pc", "pass"))
+	prvRaw, pubRaw, err := secure.GenerateKeyFile(false, "pc", "pass")
+	g.E(err)
 
-	pub, err := secure.SSHPubKey(g.Read(p + secure.PUB_KEY_EXT).Bytes())
+	pub, err := secure.SSHPubKey(pubRaw)
 	g.E(err)
 	g.Is(pub, ed25519.PublicKey{})
 
-	prv, err := secure.SSHPrvKey(g.Read(p).Bytes(), "pass")
+	prv, err := secure.SSHPrvKey(prvRaw, "pass")
 	g.E(err)
 	g.Is(prv, ed25519.PrivateKey{})
 }
 
 func TestGenerateDeterministicKeyFile(t *testing.T) {
 	g := got.T(t)
 
-	g.MkdirAll(0, "tmp")
-
-	p := "tmp/id_ed25519_deterministic"
-
-	g.E(secure.GenerateKeyFile(true, p, "pc", "pass"))
-
-	pub := g.Read(p + secure.PUB_KEY_EXT).Bytes()
-
-	g.E(secure.GenerateKeyFile(true, p, "pc", "pass"))
+	_, pub, err := secure.GenerateKeyFile(true, "pc", "pass")
+	g.E(err)
 
-	prv := g.Read(p).Bytes()
+	prv, _, err := secure.GenerateKeyFile(true, "pc", "pass")
+	g.E(err)
 
 	yes, err := secure.Belongs(pub, prv, "pass")
 	g.E(err)

lib/secure/utils.go

@@ -10,7 +10,6 @@ import (
 	"encoding/pem"
 	"fmt"
 	"math/big"
-	"os"
 
 	"golang.org/x/crypto/argon2"
 	"golang.org/x/crypto/hkdf"
@@ -83,53 +82,50 @@ func bytesToKeys(private []byte, passphrase string, publics [][]byte) (crypto.Pr
 	return prv, pubs, nil
 }
 
-// GenerateKeyFile generates a new ed25519 ssh key pair.
+// GenerateKeyFile generates a new ed25519 ssh key pair. The first return value is the private key in PEM format,
+// the second return value is the public key in ssh authorized_key format.
 // If deterministic is true, the key will be generated based on the passphrase itself,
 // so the same passphrase will always generate the same key, this is useful if you don't want to backup the key,
 // but it's less secure, you must use a strong passphrase.
-func GenerateKeyFile(deterministic bool, privateKeyPath, comment, passphrase string) error {
+func GenerateKeyFile(deterministic bool, comment, passphrase string) ([]byte, []byte, error) {
 	var prvKeyPem *pem.Block
 
 	seed := rand.Reader
 
 	if passphrase != "" && deterministic {
 		salt := sha256.Sum256([]byte(passphrase))
-		derivedKey := argon2.IDKey([]byte(passphrase), salt[:], 128, 64*1024, 4, 32)
+		derivedKey := argon2.IDKey([]byte(passphrase), salt[:], 1, 64*1024, 4, 32)
 		seed = hkdf.New(sha256.New, derivedKey, nil, nil)
 	}
 
 	publicKey, privateKey, err := ed25519.GenerateKey(seed)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
 
 	sshPubKey, err := ssh.NewPublicKey(publicKey)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
 
 	pubKeyString := fmt.Sprintf("%s %s %s\n",
 		sshPubKey.Type(),
 		base64.StdEncoding.EncodeToString(sshPubKey.Marshal()),
 		comment,
 	)
-	err = os.WriteFile(privateKeyPath+PUB_KEY_EXT, []byte(pubKeyString), 0o644)
-	if err != nil {
-		return err
-	}
 
 	if passphrase != "" {
 		prvKeyPem, err = ssh.MarshalPrivateKeyWithPassphrase(privateKey, comment, []byte(passphrase))
 	} else {
 		prvKeyPem, err = ssh.MarshalPrivateKey(privateKey, comment)
 	}
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
 
 	prvKeyBytes := pem.EncodeToMemory(prvKeyPem)
 
-	return os.WriteFile(privateKeyPath, prvKeyBytes, 0o600)
+	return prvKeyBytes, []byte(pubKeyString), nil
 }
 
 func bigIntToBytes(a *big.Int, padding int) []byte {

lib/whisper.go

@@ -12,7 +12,7 @@ import (
 )
 
 const (
-	APIVersion        = "v0.8.14"
+	APIVersion        = "v0.9.0"
 	WireFormatVersion = byte(8)
 )
 

main.go

@@ -315,7 +315,17 @@ func genKey(path string) {
 	comment := readLine("Enter the comment for it: ")
 	deterministic := readLine("Enter yes for deterministic key: ") == "yes"
 
-	err := secure.GenerateKeyFile(deterministic, path, comment, pass)
+	prv, pub, err := secure.GenerateKeyFile(deterministic, comment, pass)
+	if err != nil {
+		exit(err)
+	}
+
+	err = os.WriteFile(path, prv, 0o600)
+	if err != nil {
+		exit(err)
+	}
+
+	err = os.WriteFile(path+secure.PUB_KEY_EXT, pub, 0o644)
 	if err != nil {
 		exit(err)
 	}