Update the export / import calls

gnolang · Dec 6, 2024 · 41e8945 · 41e8945
1 parent faf6be2
commit 41e8945
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 30 deletions.
diff --git a/tm2/pkg/crypto/keys/client/export.go b/tm2/pkg/crypto/keys/client/export.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/gnolang/gno/tm2/pkg/commands"
 	"github.com/gnolang/gno/tm2/pkg/crypto/keys"
+	"github.com/gnolang/gno/tm2/pkg/crypto/keys/armor"
 )
 
 type ExportCfg struct {
@@ -88,25 +89,25 @@ func execExport(cfg *ExportCfg, io commands.IO) error {
 		)
 	}
 
-	var (
-		armor     string
-		exportErr error
-	)
+	var keyArmor string
 
 	if cfg.Unsafe {
-		// Generate the unencrypted armor
-		armor, exportErr = kb.ExportPrivKeyUnsafe(
-			cfg.NameOrBech32,
-			decryptPassword,
-		)
-
 		privk, err := kb.ExportPrivKey(cfg.NameOrBech32, decryptPassword)
 		if err != nil {
-			panic(err)
+			return fmt.Errorf("unable to export private key, %w", err)
 		}
 
 		io.Printf("privk:\n%x\n", privk.Bytes())
+
+		// Generate the private key armor
+		keyArmor = armor.ArmorPrivateKey(privk)
 	} else {
+		// Generate the encrypted armor
+		privk, err := kb.ExportPrivKey(cfg.NameOrBech32, decryptPassword)
+		if err != nil {
+			return fmt.Errorf("unable to export private key, %w", err)
+		}
+
 		// Get the armor encrypt password
 		encryptPassword, err := io.GetCheckPassword(
 			[2]string{
@@ -122,25 +123,13 @@ func execExport(cfg *ExportCfg, io commands.IO) error {
 			)
 		}
 
-		// Generate the encrypted armor
-		armor, exportErr = kb.ExportPrivKey(
-			cfg.NameOrBech32,
-			decryptPassword,
-			encryptPassword,
-		)
-	}
-
-	if exportErr != nil {
-		return fmt.Errorf(
-			"unable to export the private key, %w",
-			exportErr,
-		)
+		keyArmor = armor.EncryptArmorPrivKey(privk, encryptPassword)
 	}
 
 	// Write the armor to disk
 	if err := os.WriteFile(
 		cfg.OutputPath,
-		[]byte(armor),
+		[]byte(keyArmor),
 		0o644,
 	); err != nil {
 		return fmt.Errorf(

diff --git a/tm2/pkg/crypto/keys/client/import.go b/tm2/pkg/crypto/keys/client/import.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/gnolang/gno/tm2/pkg/commands"
 	"github.com/gnolang/gno/tm2/pkg/crypto/keys"
+	"github.com/gnolang/gno/tm2/pkg/crypto/keys/armor"
 )
 
 type ImportCfg struct {
@@ -77,7 +78,7 @@ func execImport(cfg *ImportCfg, io commands.IO) error {
 	}
 
 	// Read the raw encrypted armor
-	armor, err := os.ReadFile(cfg.ArmorPath)
+	keyArmor, err := os.ReadFile(cfg.ArmorPath)
 	if err != nil {
 		return fmt.Errorf(
 			"unable to read armor from path %s, %w",
@@ -121,10 +122,16 @@ func execImport(cfg *ImportCfg, io commands.IO) error {
 	}
 
 	if cfg.Unsafe {
+		// Un-armor the private key
+		privKey, err := armor.UnarmorPrivateKey(string(keyArmor))
+		if err != nil {
+			return fmt.Errorf("unable to unarmor private key, %w", err)
+		}
+
 		// Import the unencrypted private key
-		if err := kb.ImportPrivKeyUnsafe(
+		if err := kb.ImportPrivKey(
 			cfg.KeyName,
-			string(armor),
+			privKey,
 			encryptPassword,
 		); err != nil {
 			return fmt.Errorf(
@@ -133,11 +140,16 @@ func execImport(cfg *ImportCfg, io commands.IO) error {
 			)
 		}
 	} else {
+		// Decrypt the armor
+		privKey, err := armor.UnarmorDecryptPrivKey(string(keyArmor), decryptPassword)
+		if err != nil {
+			return fmt.Errorf("unable to decrypt private key armor, %w", err)
+		}
+
 		// Import the encrypted private key
 		if err := kb.ImportPrivKey(
 			cfg.KeyName,
-			string(armor),
-			decryptPassword,
+			privKey,
 			encryptPassword,
 		); err != nil {
 			return fmt.Errorf(

diff --git a/tm2/pkg/crypto/keys/client/import_test.go b/tm2/pkg/crypto/keys/client/import_test.go
@@ -3,12 +3,14 @@ package client
 import (
 	"fmt"
 	"io"
+	"os"
 	"strings"
 	"testing"
 
 	"github.com/gnolang/gno/tm2/pkg/commands"
 	"github.com/gnolang/gno/tm2/pkg/testutils"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type testImportKeyOpts struct {
@@ -154,6 +156,8 @@ func TestImport_ImportKey(t *testing.T) {
 }
 
 func TestImport_ImportKeyWithEmptyName(t *testing.T) {
+	t.Parallel()
+
 	// Generate a temporary key-base directory
 	_, kbHome := newTestKeybase(t)
 	err := importKey(
@@ -168,3 +172,74 @@ func TestImport_ImportKeyWithEmptyName(t *testing.T) {
 	assert.Error(t, err)
 	assert.EqualError(t, err, "name shouldn't be empty")
 }
+
+func TestImport_ImportKeyInvalidArmor(t *testing.T) {
+	t.Parallel()
+
+	_, kbHome := newTestKeybase(t)
+
+	armorFile, err := os.CreateTemp("", "armor.key")
+	require.NoError(t, err)
+
+	defer os.Remove(armorFile.Name())
+
+	// Write invalid armor
+	_, err = armorFile.Write([]byte("totally valid tendermint armor"))
+	require.NoError(t, err)
+
+	err = importKey(
+		testImportKeyOpts{
+			testCmdKeyOptsBase: testCmdKeyOptsBase{
+				kbHome:  kbHome,
+				keyName: "key-name",
+				unsafe:  true, // expect an unencrypted private key armor
+			},
+			armorPath: armorFile.Name(),
+		},
+		strings.NewReader(
+			fmt.Sprintf(
+				"%s\n%s\n",
+				"",
+				"",
+			),
+		),
+	)
+
+	assert.ErrorContains(t, err, "unable to unarmor private key")
+}
+
+func TestImport_ImportKeyInvalidPKArmor(t *testing.T) {
+	t.Parallel()
+
+	_, kbHome := newTestKeybase(t)
+
+	armorFile, err := os.CreateTemp("", "armor.key")
+	require.NoError(t, err)
+
+	defer os.Remove(armorFile.Name())
+
+	// Write invalid armor
+	_, err = armorFile.Write([]byte("totally valid tendermint armor"))
+	require.NoError(t, err)
+
+	err = importKey(
+		testImportKeyOpts{
+			testCmdKeyOptsBase: testCmdKeyOptsBase{
+				kbHome:  kbHome,
+				keyName: "key-name",
+				unsafe:  false, // expect an encrypted private key armor
+			},
+			armorPath: armorFile.Name(),
+		},
+		strings.NewReader(
+			fmt.Sprintf(
+				"%s\n%s\n%s\n",
+				"",
+				"",
+				"",
+			),
+		),
+	)
+
+	assert.ErrorContains(t, err, "unable to decrypt private key armor")
+}