From 094534e31ae722bf96485fdeaf9a6559b9a0127e Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 18 May 2022 13:47:25 +0300
Subject: [PATCH] [#1379] neofs-cli: Add `key.GetOrGenerate` helper

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 cmd/neofs-cli/internal/key/raw.go | 15 +++++++++++++++
 cmd/neofs-cli/modules/root.go     | 14 +-------------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/cmd/neofs-cli/internal/key/raw.go b/cmd/neofs-cli/internal/key/raw.go
index 15ee4b255..7e2a921b2 100644
--- a/cmd/neofs-cli/internal/key/raw.go
+++ b/cmd/neofs-cli/internal/key/raw.go
@@ -2,6 +2,7 @@ package key
 
 import (
 	"crypto/ecdsa"
+	"errors"
 	"fmt"
 	"os"
 
@@ -11,6 +12,8 @@ import (
 	"github.com/spf13/viper"
 )
 
+var errCantGenerateKey = errors.New("can't generate new private key")
+
 // Get returns private key from the following sources:
 // 1. WIF
 // 2. Raw binary key
@@ -43,6 +46,18 @@ func Get() (*ecdsa.PrivateKey, error) {
 	return nil, ErrInvalidKey
 }
 
+// GetOrGenerate is similar to get but generates a new key if commonflags.GenerateKey is set.
+func GetOrGenerate() (*ecdsa.PrivateKey, error) {
+	if viper.GetBool(commonflags.GenerateKey) {
+		priv, err := keys.NewPrivateKey()
+		if err != nil {
+			return nil, fmt.Errorf("%w: %v", errCantGenerateKey, err)
+		}
+		return &priv.PrivateKey, nil
+	}
+	return Get()
+}
+
 func getKeyFromFile(keyPath string) (*ecdsa.PrivateKey, error) {
 	data, err := os.ReadFile(keyPath)
 	if err != nil {
diff --git a/cmd/neofs-cli/modules/root.go b/cmd/neofs-cli/modules/root.go
index 7bbc45c9a..e8e827cec 100644
--- a/cmd/neofs-cli/modules/root.go
+++ b/cmd/neofs-cli/modules/root.go
@@ -9,7 +9,6 @@ import (
 	"strings"
 
 	"github.com/mitchellh/go-homedir"
-	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	internalclient "github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/client"
 	"github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/commonflags"
 	"github.com/nspcc-dev/neofs-node/cmd/neofs-cli/internal/key"
@@ -62,10 +61,6 @@ and much more!`,
 	Run: entryPoint,
 }
 
-var (
-	errCantGenerateKey = errors.New("can't generate new private key")
-)
-
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
@@ -140,14 +135,7 @@ func initConfig() {
 
 // getKey returns private key that was provided in global arguments.
 func getKey() (*ecdsa.PrivateKey, error) {
-	if viper.GetBool(commonflags.GenerateKey) {
-		priv, err := keys.NewPrivateKey()
-		if err != nil {
-			return nil, errCantGenerateKey
-		}
-		return &priv.PrivateKey, nil
-	}
-	return getKeyNoGenerate()
+	return key.GetOrGenerate()
 }
 
 func getKeyNoGenerate() (*ecdsa.PrivateKey, error) {