Merge pull request #757 from nspcc-dev/feature/cli

cli: implement `wallet remove-account`

cli/wallet/nep5.go

@@ -66,6 +66,20 @@ func newNEP5Commands() []cli.Command {
 				},
 			},
 		},
+		{
+			Name:      "remove",
+			Usage:     "remove NEP5 token from the wallet",
+			UsageText: "remove --path <path> <hash-or-name>",
+			Action:    removeNEP5Token,
+			Flags: []cli.Flag{
+				walletPathFlag,
+				cli.StringFlag{
+					Name:  "token",
+					Usage: "Token name or hash",
+				},
+				forceFlag,
+			},
+		},
 		{
 			Name:      "transfer",
 			Usage:     "transfer NEP5 tokens",
@@ -260,6 +274,34 @@ func printNEP5Info(ctx *cli.Context) error {
 	return nil
 }
 
+func removeNEP5Token(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	defer wall.Close()
+
+	name := ctx.Args().First()
+	if name == "" {
+		return cli.NewExitError("token must be specified", 1)
+	}
+	token, err := getMatchingToken(wall, name)
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	if !ctx.Bool("force") {
+		if ok := askForConsent(); !ok {
+			return nil
+		}
+	}
+	if err := wall.RemoveToken(token.Hash); err != nil {
+		return cli.NewExitError(fmt.Errorf("can't remove token: %v", err), 1)
+	} else if err := wall.Save(); err != nil {
+		return cli.NewExitError(fmt.Errorf("error while saving wallet: %v", err), 1)
+	}
+	return nil
+}
+
 func transferNEP5(ctx *cli.Context) error {
 	wall, err := openWallet(ctx.String("path"))
 	if err != nil {

cli/wallet/wallet.go

@@ -3,6 +3,7 @@ package wallet
 import (
 	"bufio"
 	"context"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -67,6 +68,10 @@ var (
 		Name:  "to",
 		Usage: "Address to send an asset to",
 	}
+	forceFlag = cli.BoolFlag{
+		Name:  "force",
+		Usage: "Do not ask for a confirmation",
+	}
 )
 
 // NewCommands returns 'wallet' command.
@@ -90,7 +95,7 @@ func NewCommands() []cli.Command {
 				},
 			},
 			{
-				Name:   "create",
+				Name:   "init",
 				Usage:  "create a new wallet",
 				Action: createWallet,
 				Flags: []cli.Flag{
@@ -102,7 +107,7 @@ func NewCommands() []cli.Command {
 				},
 			},
 			{
-				Name:   "create-account",
+				Name:   "create",
 				Usage:  "add an account to the existing wallet",
 				Action: addAccount,
 				Flags: []cli.Flag{
@@ -139,6 +144,10 @@ func NewCommands() []cli.Command {
 						Name:  "name, n",
 						Usage: "Optional account name",
 					},
+					cli.StringFlag{
+						Name:  "contract",
+						Usage: "Verification script for custom contracts",
+					},
 				},
 			},
 			{
@@ -160,6 +169,16 @@ func NewCommands() []cli.Command {
 					},
 				},
 			},
+			{
+				Name:      "remove",
+				Usage:     "remove an account from the wallet",
+				UsageText: "remove --path <path> [--force] <addr>",
+				Action:    removeAccount,
+				Flags: []cli.Flag{
+					walletPathFlag,
+					forceFlag,
+				},
+			},
 			{
 				Name:  "transfer",
 				Usage: "transfer NEO/GAS",
@@ -389,6 +408,14 @@ func importWallet(ctx *cli.Context) error {
 		return cli.NewExitError(err, 1)
 	}
 
+	if ctrFlag := ctx.String("contract"); ctrFlag != "" {
+		ctr, err := hex.DecodeString(ctrFlag)
+		if err != nil {
+			return cli.NewExitError("invalid contract", 1)
+		}
+		acc.Contract.Script = ctr
+	}
+
 	acc.Label = ctx.String("name")
 	if err := addAccountAndSave(wall, acc); err != nil {
 		return cli.NewExitError(err, 1)
@@ -397,6 +424,52 @@ func importWallet(ctx *cli.Context) error {
 	return nil
 }
 
+func removeAccount(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	defer wall.Close()
+
+	addrArg := ctx.Args().First()
+	addr, err := address.StringToUint160(addrArg)
+	if err != nil {
+		return cli.NewExitError("valid address must be provided", 1)
+	}
+	acc := wall.GetAccount(addr)
+	if acc == nil {
+		return cli.NewExitError("account wasn't found", 1)
+	}
+
+	if !ctx.Bool("force") {
+		fmt.Printf("Account %s will be removed. This action is irreversible.\n", addrArg)
+		if ok := askForConsent(); !ok {
+			return nil
+		}
+	}
+
+	if err := wall.RemoveAccount(acc.Address); err != nil {
+		return cli.NewExitError(fmt.Errorf("error on remove: %v", err), 1)
+	} else if err := wall.Save(); err != nil {
+		return cli.NewExitError(fmt.Errorf("error while saving wallet: %v", err), 1)
+	}
+	return nil
+}
+
+func askForConsent() bool {
+	fmt.Print("Are you sure? [y/N]: ")
+	reader := bufio.NewReader(os.Stdin)
+	response, err := reader.ReadString('\n')
+	if err == nil {
+		response = strings.ToLower(strings.TrimSpace(response))
+		if response == "y" || response == "yes" {
+			return true
+		}
+	}
+	fmt.Println("Cancelled.")
+	return false
+}
+
 func transferAsset(ctx *cli.Context) error {
 	wall, err := openWallet(ctx.String("path"))
 	if err != nil {

pkg/wallet/wallet.go

@@ -2,6 +2,7 @@ package wallet
 
 import (
 	"encoding/json"
+	"errors"
 	"io"
 	"os"
 
@@ -102,11 +103,36 @@ func (w *Wallet) AddAccount(acc *Account) {
 	w.Accounts = append(w.Accounts, acc)
 }
 
+// RemoveAccount removes an Account with the specified addr
+// from the wallet.
+func (w *Wallet) RemoveAccount(addr string) error {
+	for i, acc := range w.Accounts {
+		if acc.Address == addr {
+			copy(w.Accounts[i:], w.Accounts[i+1:])
+			w.Accounts = w.Accounts[:len(w.Accounts)-1]
+			return nil
+		}
+	}
+	return errors.New("account wasn't found")
+}
+
 // AddToken adds new token to a wallet.
 func (w *Wallet) AddToken(tok *Token) {
 	w.Extra.Tokens = append(w.Extra.Tokens, tok)
 }
 
+// RemoveToken removes token with the specified hash from the wallet.
+func (w *Wallet) RemoveToken(h util.Uint160) error {
+	for i, tok := range w.Extra.Tokens {
+		if tok.Hash.Equals(h) {
+			copy(w.Extra.Tokens[i:], w.Extra.Tokens[i+1:])
+			w.Extra.Tokens = w.Extra.Tokens[:len(w.Extra.Tokens)-1]
+			return nil
+		}
+	}
+	return errors.New("token wasn't found")
+}
+
 // Path returns the location of the wallet on the filesystem.
 func (w *Wallet) Path() string {
 	return w.path

pkg/wallet/wallet_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
+	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -48,7 +49,7 @@ func TestAddAccount(t *testing.T) {
 		privateKey:   nil,
 		publicKey:    nil,
 		wif:          "",
-		Address:      "",
+		Address:      "real",
 		EncryptedWIF: "",
 		Label:        "",
 		Contract:     nil,
@@ -57,6 +58,11 @@ func TestAddAccount(t *testing.T) {
 	})
 	accounts := wallet.Accounts
 	require.Len(t, accounts, 1)
+
+	require.Error(t, wallet.RemoveAccount("abc"))
+	require.Len(t, wallet.Accounts, 1)
+	require.NoError(t, wallet.RemoveAccount("real"))
+	require.Len(t, wallet.Accounts, 0)
 }
 
 func TestPath(t *testing.T) {
@@ -133,6 +139,18 @@ func removeWallet(t *testing.T, walletPath string) {
 	require.NoError(t, err)
 }
 
+func TestWallet_AddToken(t *testing.T) {
+	w := checkWalletConstructor(t)
+	tok := NewToken(util.Uint160{1, 2, 3}, "Rubl", "RUB", 2)
+	require.Equal(t, 0, len(w.Extra.Tokens))
+	w.AddToken(tok)
+	require.Equal(t, 1, len(w.Extra.Tokens))
+	require.Error(t, w.RemoveToken(util.Uint160{4, 5, 6}))
+	require.Equal(t, 1, len(w.Extra.Tokens))
+	require.NoError(t, w.RemoveToken(tok.Hash))
+	require.Equal(t, 0, len(w.Extra.Tokens))
+}
+
 func TestWallet_GetAccount(t *testing.T) {
 	wallet := checkWalletConstructor(t)
 	accounts := []*Account{