From e8d552cf75b08c3bcd0b5b83b8ece9fa76a82b14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matev=C5=BE=20Jekovec?= <matevz@oasisprotocol.org>
Date: Thu, 19 Oct 2023 12:16:57 +0200
Subject: [PATCH] feat(cmd/wallet): Import key from PEM file

---
 cmd/wallet/create.go      | 14 +++++--
 cmd/wallet/import.go      |  9 +---
 cmd/wallet/import_file.go | 88 +++++++++++++++++++++++++++++++++++++++
 cmd/wallet/wallet.go      |  1 +
 4 files changed, 101 insertions(+), 11 deletions(-)
 create mode 100644 cmd/wallet/import_file.go

diff --git a/cmd/wallet/create.go b/cmd/wallet/create.go
index db709406..5fdfcf1a 100644
--- a/cmd/wallet/create.go
+++ b/cmd/wallet/create.go
@@ -22,6 +22,8 @@ var createCmd = &cobra.Command{
 		cfg := config.Global()
 		name := args[0]
 
+		checkAccountExists(cfg, name)
+
 		af, err := wallet.Load(accKind)
 		cobra.CheckErr(err)
 
@@ -37,9 +39,6 @@ var createCmd = &cobra.Command{
 		err = accCfg.SetConfigFromFlags()
 		cobra.CheckErr(err)
 
-		if _, exists := cfg.AddressBook.All[name]; exists {
-			cobra.CheckErr(fmt.Errorf("address named '%s' already exists in address book", name))
-		}
 		err = cfg.Wallet.Create(name, passphrase, accCfg)
 		cobra.CheckErr(err)
 
@@ -48,6 +47,15 @@ var createCmd = &cobra.Command{
 	},
 }
 
+func checkAccountExists(cfg *config.Config, name string) {
+	if _, exists := cfg.Wallet.All[name]; exists {
+		cobra.CheckErr(fmt.Errorf("account '%s' already exists in the wallet", name))
+	}
+	if _, exists := cfg.AddressBook.All[name]; exists {
+		cobra.CheckErr(fmt.Errorf("address named '%s' already exists in the address book", name))
+	}
+}
+
 func init() {
 	flags := flag.NewFlagSet("", flag.ContinueOnError)
 	kinds := make([]string, 0, len(wallet.AvailableKinds()))
diff --git a/cmd/wallet/import.go b/cmd/wallet/import.go
index 1184236a..19ec77d4 100644
--- a/cmd/wallet/import.go
+++ b/cmd/wallet/import.go
@@ -1,8 +1,6 @@
 package wallet
 
 import (
-	"fmt"
-
 	"github.com/AlecAivazis/survey/v2"
 	"github.com/spf13/cobra"
 
@@ -20,12 +18,7 @@ var importCmd = &cobra.Command{
 		cfg := config.Global()
 		name := args[0]
 
-		if _, exists := cfg.Wallet.All[name]; exists {
-			cobra.CheckErr(fmt.Errorf("account '%s' already exists in the wallet", name))
-		}
-		if _, exists := cfg.AddressBook.All[name]; exists {
-			cobra.CheckErr(fmt.Errorf("address named '%s' already exists in the address book", name))
-		}
+		checkAccountExists(cfg, name)
 
 		// NOTE: We only support importing into the file-based wallet for now.
 		af, err := wallet.Load(walletFile.Kind)
diff --git a/cmd/wallet/import_file.go b/cmd/wallet/import_file.go
new file mode 100644
index 00000000..28d7ce6e
--- /dev/null
+++ b/cmd/wallet/import_file.go
@@ -0,0 +1,88 @@
+package wallet
+
+import (
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/pem"
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+
+	"github.com/oasisprotocol/cli/cmd/common"
+	"github.com/oasisprotocol/cli/config"
+	"github.com/oasisprotocol/cli/wallet"
+	walletFile "github.com/oasisprotocol/cli/wallet/file"
+)
+
+var importFileCmd = &cobra.Command{
+	Use:   "import-file <file_name> <account_name>",
+	Short: "Import an existing account from file",
+	Long:  "Import the private key from an existing PEM file",
+	Args:  cobra.ExactArgs(2),
+	Run: func(cmd *cobra.Command, args []string) {
+		cfg := config.Global()
+		filename := args[0]
+		name := args[1]
+
+		checkAccountExists(cfg, name)
+
+		rawFile, err := os.ReadFile(filename)
+		cobra.CheckErr(err)
+
+		block, _ := pem.Decode(rawFile)
+		if block == nil {
+			cobra.CheckErr(fmt.Errorf("failed to decode PEM file"))
+		}
+
+		algorithm, err := detectAlgorithm(block.Type)
+		cobra.CheckErr(err)
+
+		// Ask for passphrase.
+		passphrase := common.AskNewPassphrase()
+
+		accCfg := &config.Account{
+			Kind: walletFile.Kind,
+			Config: map[string]interface{}{
+				"algorithm": algorithm,
+			},
+		}
+
+		src := &wallet.ImportSource{
+			Kind: wallet.ImportKindPrivateKey,
+			Data: encodeKeyData(algorithm, block.Bytes),
+		}
+
+		err = cfg.Wallet.Import(name, passphrase, accCfg, src)
+		cobra.CheckErr(err)
+
+		err = cfg.Save()
+		cobra.CheckErr(err)
+	},
+}
+
+// detectAlgorithm detects the key type based on the PEM type.
+func detectAlgorithm(pemType string) (string, error) {
+	switch pemType {
+	case "ED25519 PRIVATE KEY":
+		return wallet.AlgorithmEd25519Raw, nil
+	case "EC PRIVATE KEY":
+		return wallet.AlgorithmSecp256k1Raw, nil
+	case "SR25519 PRIVATE KEY":
+		return wallet.AlgorithmSr25519Raw, nil
+	}
+
+	return "", fmt.Errorf("unsupported PEM type: %s", pemType)
+}
+
+// encodeKeyData re-encodes the key in raw bytes back to the user-readable string for import.
+func encodeKeyData(algorithm string, rawKey []byte) string {
+	switch algorithm {
+	case wallet.AlgorithmEd25519Raw, wallet.AlgorithmSr25519Raw:
+		return base64.StdEncoding.EncodeToString(rawKey)
+	case wallet.AlgorithmSecp256k1Raw:
+		return hex.EncodeToString(rawKey)
+	}
+
+	return ""
+}
diff --git a/cmd/wallet/wallet.go b/cmd/wallet/wallet.go
index 43f80229..8594f0fd 100644
--- a/cmd/wallet/wallet.go
+++ b/cmd/wallet/wallet.go
@@ -53,6 +53,7 @@ func init() {
 	Cmd.AddCommand(renameCmd)
 	Cmd.AddCommand(setDefaultCmd)
 	Cmd.AddCommand(importCmd)
+	Cmd.AddCommand(importFileCmd)
 	Cmd.AddCommand(exportCmd)
 	Cmd.AddCommand(remoteSignerCmd)
 }