From ced90e05108cd67c67d671f08775ebddfdecd9dd Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Fri, 16 Apr 2021 12:29:46 +0300
Subject: [PATCH] cli: use AddressFlag for `nep17 import`

---
 cli/nep17_test.go   |  8 +++++++-
 cli/wallet/nep17.go | 11 ++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/cli/nep17_test.go b/cli/nep17_test.go
index 20e5bb984..1f7b9c824 100644
--- a/cli/nep17_test.go
+++ b/cli/nep17_test.go
@@ -200,6 +200,12 @@ func TestNEP17ImportToken(t *testing.T) {
 	gasContractHash, err := e.Chain.GetNativeContractScriptHash(nativenames.Gas)
 	require.NoError(t, err)
 	e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath)
+
+	// missing token hash
+	e.RunWithError(t, "neo-go", "wallet", "nep17", "import",
+		"--rpc-endpoint", "http://"+e.RPC.Addr,
+		"--wallet", walletPath)
+
 	e.Run(t, "neo-go", "wallet", "nep17", "import",
 		"--rpc-endpoint", "http://"+e.RPC.Addr,
 		"--wallet", walletPath,
@@ -207,7 +213,7 @@ func TestNEP17ImportToken(t *testing.T) {
 	e.Run(t, "neo-go", "wallet", "nep17", "import",
 		"--rpc-endpoint", "http://"+e.RPC.Addr,
 		"--wallet", walletPath,
-		"--token", neoContractHash.StringLE())
+		"--token", address.Uint160ToString(neoContractHash)) // try address instead of sh
 
 	t.Run("Info", func(t *testing.T) {
 		checkGASInfo := func(t *testing.T) {
diff --git a/cli/wallet/nep17.go b/cli/wallet/nep17.go
index 09ba5933c..a2185946a 100644
--- a/cli/wallet/nep17.go
+++ b/cli/wallet/nep17.go
@@ -40,9 +40,9 @@ func newNEP17Commands() []cli.Command {
 	balanceFlags = append(balanceFlags, options.RPC...)
 	importFlags := []cli.Flag{
 		walletPathFlag,
-		cli.StringFlag{
+		flags.AddressFlag{
 			Name:  "token",
-			Usage: "Token contract hash in LE",
+			Usage: "Token contract address or hash in LE",
 		},
 	}
 	importFlags = append(importFlags, options.RPC...)
@@ -257,10 +257,11 @@ func importNEP17Token(ctx *cli.Context) error {
 	}
 	defer wall.Close()
 
-	tokenHash, err := flags.ParseAddress(ctx.String("token"))
-	if err != nil {
-		return cli.NewExitError(fmt.Errorf("invalid token contract hash: %w", err), 1)
+	tokenHashFlag := ctx.Generic("token").(*flags.Address)
+	if !tokenHashFlag.IsSet {
+		return cli.NewExitError("token contract hash was not set", 1)
 	}
+	tokenHash := tokenHashFlag.Uint160()
 
 	for _, t := range wall.Extra.Tokens {
 		if t.Hash.Equals(tokenHash) {