[#1424] neofs-cli: Fail immediately if a key can't be fetched

If the key can't be fetched, an error is always returned, so it makes
sense to fail the whole command inside of a `key.Get*()`.

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>

cmd/neofs-cli/internal/key/key_test.go

@@ -0,0 +1,132 @@
+package key
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/cli/input"
+	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/wallet"
+	"github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/commonflags"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/term"
+)
+
+func Test_getOrGenerate(t *testing.T) {
+	dir := t.TempDir()
+
+	wallPath := filepath.Join(dir, "wallet.json")
+	w, err := wallet.NewWallet(wallPath)
+	require.NoError(t, err)
+
+	acc1, err := wallet.NewAccount()
+	require.NoError(t, err)
+	require.NoError(t, acc1.Encrypt("pass", keys.NEP2ScryptParams()))
+	w.AddAccount(acc1)
+
+	acc2, err := wallet.NewAccount()
+	require.NoError(t, err)
+	require.NoError(t, acc2.Encrypt("pass", keys.NEP2ScryptParams()))
+	acc2.Default = true
+	w.AddAccount(acc2)
+	require.NoError(t, w.Save())
+	w.Close()
+
+	keyPath := filepath.Join(dir, "binary.key")
+	rawKey, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+	require.NoError(t, os.WriteFile(keyPath, rawKey.Bytes(), os.ModePerm))
+
+	wifKey, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+
+	nep2Key, err := keys.NewPrivateKey()
+	require.NoError(t, err)
+	nep2, err := keys.NEP2Encrypt(nep2Key, "pass", keys.NEP2ScryptParams())
+	require.NoError(t, err)
+
+	in := bytes.NewBuffer(nil)
+	input.Terminal = term.NewTerminal(input.ReadWriter{
+		Reader: in,
+		Writer: io.Discard,
+	}, "")
+
+	checkKeyError(t, filepath.Join(dir, "badfile"), ErrInvalidKey)
+
+	t.Run("wallet", func(t *testing.T) {
+		checkKeyError(t, wallPath, ErrInvalidPassword)
+
+		in.WriteString("invalid\r")
+		checkKeyError(t, wallPath, ErrInvalidPassword)
+
+		in.WriteString("pass\r")
+		checkKey(t, wallPath, acc2.PrivateKey()) // default account
+
+		viper.Set(commonflags.Account, acc1.Address)
+		in.WriteString("pass\r")
+		checkKey(t, wallPath, acc1.PrivateKey())
+
+		viper.Set(commonflags.Account, "not an address")
+		checkKeyError(t, wallPath, ErrInvalidAddress)
+
+		acc, err := wallet.NewAccount()
+		require.NoError(t, err)
+		viper.Set(commonflags.Account, acc.Address)
+		checkKeyError(t, wallPath, ErrInvalidAddress)
+	})
+
+	t.Run("WIF", func(t *testing.T) {
+		checkKey(t, wifKey.WIF(), wifKey)
+	})
+
+	t.Run("NEP-2", func(t *testing.T) {
+		checkKeyError(t, nep2, ErrInvalidPassword)
+
+		in.WriteString("invalid\r")
+		checkKeyError(t, nep2, ErrInvalidPassword)
+
+		in.WriteString("pass\r")
+		checkKey(t, nep2, nep2Key)
+
+		t.Run("password from config", func(t *testing.T) {
+			viper.Set("password", "invalid")
+			in.WriteString("pass\r")
+			checkKeyError(t, nep2, ErrInvalidPassword)
+
+			viper.Set("password", "pass")
+			in.WriteString("invalid\r")
+			checkKey(t, nep2, nep2Key)
+		})
+	})
+
+	t.Run("raw key", func(t *testing.T) {
+		checkKey(t, keyPath, rawKey)
+	})
+
+	t.Run("generate", func(t *testing.T) {
+		viper.Set(commonflags.GenerateKey, true)
+		actual, err := getOrGenerate()
+		require.NoError(t, err)
+		require.NotNil(t, actual)
+		for _, p := range []*keys.PrivateKey{nep2Key, rawKey, wifKey, acc1.PrivateKey(), acc2.PrivateKey()} {
+			require.NotEqual(t, p, actual, "expected new key to be generated")
+		}
+	})
+}
+
+func checkKeyError(t *testing.T, desc string, err error) {
+	viper.Set(commonflags.WalletPath, desc)
+	_, actualErr := getOrGenerate()
+	require.ErrorIs(t, actualErr, err)
+}
+
+func checkKey(t *testing.T, desc string, expected *keys.PrivateKey) {
+	viper.Set(commonflags.WalletPath, desc)
+	actual, err := getOrGenerate()
+	require.NoError(t, err)
+	require.Equal(t, &expected.PrivateKey, actual)
+}

cmd/neofs-cli/internal/key/raw.go

@@ -8,7 +8,9 @@ import (
 
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
+	"github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/common"
 	"github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/commonflags"
+	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
 
@@ -22,7 +24,13 @@ var errCantGenerateKey = errors.New("can't generate new private key")
 // Ideally we want to touch file-system on the last step.
 // However, asking for NEP-2 password seems to be confusing if we provide a wallet.
 // This function assumes that all flags were bind to viper in a `PersistentPreRun`.
-func Get() (*ecdsa.PrivateKey, error) {
+func Get(cmd *cobra.Command) *ecdsa.PrivateKey {
+	pk, err := get()
+	common.ExitOnErr(cmd, "can't fetch private key: %w", err)
+	return pk
+}
+
+func get() (*ecdsa.PrivateKey, error) {
 	keyDesc := viper.GetString(commonflags.WalletPath)
 	priv, err := keys.NewPrivateKeyFromWIF(keyDesc)
 	if err == nil {
@@ -47,7 +55,13 @@ func Get() (*ecdsa.PrivateKey, error) {
 }
 
 // GetOrGenerate is similar to get but generates a new key if commonflags.GenerateKey is set.
-func GetOrGenerate() (*ecdsa.PrivateKey, error) {
+func GetOrGenerate(cmd *cobra.Command) *ecdsa.PrivateKey {
+	pk, err := getOrGenerate()
+	common.ExitOnErr(cmd, "can't fetch private key: %w", err)
+	return pk
+}
+
+func getOrGenerate() (*ecdsa.PrivateKey, error) {
 	if viper.GetBool(commonflags.GenerateKey) {
 		priv, err := keys.NewPrivateKey()
 		if err != nil {
@@ -55,7 +69,7 @@ func GetOrGenerate() (*ecdsa.PrivateKey, error) {
 		}
 		return &priv.PrivateKey, nil
 	}
-	return Get()
+	return get()
 }
 
 func getKeyFromFile(keyPath string) (*ecdsa.PrivateKey, error) {