wallet: make --wif optional in import-multisig command

We have a full list of public keys in the import-multisig command, so if
 we have a key in the wallet that corresponds to one of these keys
 (simple sig), just reuse it for the account automatically

Closes #3266

Signed-off-by: Ekaterina Pavlova <[email protected]>

cli/wallet/wallet.go

@@ -236,8 +236,14 @@ func NewCommands() []cli.Command {
 			{
 				Name:  "import-multisig",
 				Usage: "import multisig contract",
-				UsageText: "import-multisig -w wallet [--wallet-config path] --wif <wif> [--name <account_name>] --min <n>" +
+				UsageText: "import-multisig -w wallet [--wallet-config path] [--wif <wif>] [--name <account_name>] --min <n>" +
 					" [<pubkey1> [<pubkey2> [...]]]",
+				Description: `Imports a multisig contract with the given public keys. If --wif flag is
+	   provided, it's used to create an account with the given name (or without a name
+	   if --name flag is not provided). Otherwise, the command tries to find an account
+	   with one of the given public keys and convert it to multisig. If no account is
+	   found, an error is returned.
+`,
 				Action: importMultisig,
 				Flags: []cli.Flag{
 					walletPathFlag,
@@ -547,22 +553,118 @@ func importMultisig(ctx *cli.Context) error {
 		l := ctx.String("name")
 		label = &l
 	}
-	acc, err := newAccountFromWIF(ctx.App.Writer, ctx.String("wif"), wall.Scrypt, label, pass)
-	if err != nil {
-		return cli.NewExitError(err, 1)
+
+	var acc *wallet.Account
+
+	var matchingAccounts []*wallet.Account
+	for _, pub := range pubs {
+		match, err := findAccountsByPublicKey(wall, pub)
+		if err == nil && len(match) > 0 {
+			matchingAccounts = append(matchingAccounts, match...)
+		}
 	}
 
-	if err := acc.ConvertMultisig(m, pubs); err != nil {
-		return cli.NewExitError(err, 1)
+	if len(matchingAccounts) == 1 {
+		acc = matchingAccounts[0]
+		// Decrypt the account before converting it to multisig
+		err := decryptAccount(acc, pass, wall.Scrypt, label)
+		if err != nil {
+			return cli.NewExitError(fmt.Errorf("failed to decrypt account: %w", err), 1)
+		}
+	} else {
+		if ctx.IsSet("wif") {
+			acc, err = newAccountFromWIF(ctx.App.Writer, ctx.String("wif"), wall.Scrypt, label, pass)
+			if err != nil {
+				return cli.NewExitError(err, 1)
+			}
+		}
+	}
+	if acc != nil {
+		// Convert the account to multisig
+		if err := acc.ConvertMultisig(m, pubs); err != nil {
+			return cli.NewExitError(err, 1)
+		}
+		if len(matchingAccounts) != 1 {
+			if err := addAccountAndSave(wall, acc); err != nil {
+				return cli.NewExitError(err, 1)
+			}
+		} else {
+			wall.AddAccount(acc)
+			if err := wall.Save(); err != nil {
+				return cli.NewExitError(err, 1)
+			}
+		}
+	} else {
+		return cli.NewExitError(errors.New("none of the provided public keys correspond to an existing key in the wallet, and no WIF is provided"), 1)
 	}
+	return nil
+}
 
-	if err := addAccountAndSave(wall, acc); err != nil {
-		return cli.NewExitError(err, 1)
+// decryptAccount decrypts the given account using the provided password.
+// It returns an error if the decryption fails.
+func decryptAccount(acc *wallet.Account, pass *string, scrypt keys.ScryptParams, label *string) error {
+	var (
+		phrase, name string
+		err          error
+	)
+	// Check if the account is already decrypted
+	if acc.PrivateKey() != nil {
+		return nil
 	}
+	if pass != nil {
+		phrase = *pass
+	}
+	if label == nil {
+		name, err = readAccountName()
+		if err != nil {
+			return fmt.Errorf("failed to read account label: %w", err)
+		}
+	} else {
+		name = *label
+	}
+	if pass == nil {
+		phrase, err = input.ReadPassword(EnterPasswordPrompt)
+		if err != nil {
+			return fmt.Errorf("error reading password: %w", err)
+		}
+	}
+	// Perform the decryption
+	err = acc.Decrypt(phrase, scrypt)
+	if err != nil {
+		// If password from wallet config wasn't OK then retry with the user input,
+		// see the https://github.com/nspcc-dev/neo-go/issues/2883#issuecomment-1399923088.
+		if pass == nil {
+			return err
+		}
+		phrase, err = input.ReadPassword(EnterPasswordPrompt)
+		if err != nil {
+			return fmt.Errorf("error reading password: %w", err)
+		}
+		err = acc.Decrypt(phrase, scrypt)
+		if err != nil {
+			return err
+		}
+	}
+	acc.Label = name
 
 	return nil
 }
 
+// findAccountsByPublicKey returns all accounts in the given wallet that correspond
+// to the given public key.
+func findAccountsByPublicKey(wall *wallet.Wallet, pub *keys.PublicKey) ([]*wallet.Account, error) {
+	var accounts []*wallet.Account
+	for _, acc := range wall.Accounts {
+		if acc.ScriptHash().Equals(pub.GetScriptHash()) {
+			accounts = append(accounts, acc)
+		}
+	}
+	if len(accounts) == 0 {
+		return nil, cli.NewExitError(fmt.Errorf("no account found for public key %s", pub), 1)
+	}
+	return accounts, nil
+}
+
 func importDeployed(ctx *cli.Context) error {
 	if err := cmdargs.EnsureNone(ctx); err != nil {
 		return err

cli/wallet/wallet_test.go

@@ -487,6 +487,57 @@ func TestWalletInit(t *testing.T) {
 					hex.EncodeToString(pubs[2].Bytes()),
 					hex.EncodeToString(pubs[3].Bytes()))...)
 			})
+
+			privs, pubs = testcli.GenerateKeys(t, 3)
+			script, err = smartcontract.CreateMultiSigRedeemScript(2, pubs)
+			require.NoError(t, err)
+			// Create a wallet and import a standard account
+			e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath)
+			e.In.WriteString("standardacc\rstdpass\rstdpass\r")
+			e.Run(t, "neo-go", "wallet", "import",
+				"--wallet", walletPath,
+				"--wif", privs[0].WIF())
+			w, err = wallet.NewWalletFromFile(walletPath)
+			require.NoError(t, err)
+			actual = w.GetAccount(privs[0].GetScriptHash())
+			require.NotNil(t, actual)
+			require.NotEqual(t, actual.Contract.Script, script)
+
+			// Test when a public key of an already imported account is present
+			t.Run("existing account public key, no WIF", func(t *testing.T) {
+				e.In.WriteString("standardacc\rstdpass\rstdpass\r")
+				e.Run(t, "neo-go", "wallet", "import-multisig",
+					"--wallet", walletPath,
+					"--min", "2",
+					hex.EncodeToString(pubs[0].Bytes()), // Public key of the already imported account
+					hex.EncodeToString(pubs[1].Bytes()),
+					hex.EncodeToString(pubs[2].Bytes()))
+
+				w, err := wallet.NewWalletFromFile(walletPath)
+				require.NoError(t, err)
+				actual := w.GetAccount(hash.Hash160(script))
+				require.NotNil(t, actual)
+				require.Equal(t, actual.Contract.Script, script)
+			})
+
+			// Test when no public key of an already imported account is present, and no WIF is provided
+			t.Run("no existing account public key, no WIF", func(t *testing.T) {
+				_, pubsNew := testcli.GenerateKeys(t, 3)
+				scriptNew, err := smartcontract.CreateMultiSigRedeemScript(2, pubsNew)
+				require.NoError(t, err)
+
+				e.RunWithError(t, "neo-go", "wallet", "import-multisig",
+					"--wallet", walletPath,
+					"--min", "2",
+					hex.EncodeToString(pubsNew[0].Bytes()),
+					hex.EncodeToString(pubsNew[1].Bytes()),
+					hex.EncodeToString(pubsNew[2].Bytes()))
+
+				w, err := wallet.NewWalletFromFile(walletPath)
+				require.NoError(t, err)
+				actual := w.GetAccount(hash.Hash160(scriptNew))
+				require.Nil(t, actual)
+			})
 		})
 	})
 }