From 7cfcf072b832ff4c848c0a107efc82f38ebfdf83 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 29 Aug 2022 18:43:25 +0300
Subject: [PATCH] cli/wallet: make NEP-17 token selection logic a bit more
 robust

1. In the single token mode compare known hashes instead of names, names can
   be misleading.
2. Hardcode NEO/GAS, they are special (if not overrided by the wallet data).
---
 cli/nep17_test.go   |   1 +
 cli/wallet/nep17.go | 115 +++++++++++++++++++++++++-------------------
 docs/cli.md         |   7 +--
 3 files changed, 70 insertions(+), 53 deletions(-)

diff --git a/cli/nep17_test.go b/cli/nep17_test.go
index 2d00dfbfb..b318616d1 100644
--- a/cli/nep17_test.go
+++ b/cli/nep17_test.go
@@ -106,6 +106,7 @@ func TestNEP17Balance(t *testing.T) {
 	t.Run("Bad token", func(t *testing.T) {
 		e.Run(t, append(cmd, "--token", "kek")...)
 		e.checkNextLine(t, "^\\s*Account\\s+"+validatorAddr)
+		e.checkNextLine(t, `^\s*Can't find data for "kek" token\s*`)
 		e.checkEOF(t)
 	})
 	t.Run("Bad wallet", func(t *testing.T) {
diff --git a/cli/wallet/nep17.go b/cli/wallet/nep17.go
index 3ef164bbe..891bd0adc 100644
--- a/cli/wallet/nep17.go
+++ b/cli/wallet/nep17.go
@@ -13,6 +13,7 @@ import (
 	"github.com/nspcc-dev/neo-go/cli/input"
 	"github.com/nspcc-dev/neo-go/cli/options"
 	"github.com/nspcc-dev/neo-go/cli/paramcontext"
+	"github.com/nspcc-dev/neo-go/pkg/core/native/nativenames"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/fixedn"
@@ -20,7 +21,6 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/actor"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/gas"
-	"github.com/nspcc-dev/neo-go/pkg/rpcclient/invoker"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/neo"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/nep11"
 	"github.com/nspcc-dev/neo-go/pkg/rpcclient/nep17"
@@ -100,8 +100,21 @@ func newNEP17Commands() []cli.Command {
 			Name:      "balance",
 			Usage:     "get address balance",
 			UsageText: "balance -w wallet [--wallet-config path] --rpc-endpoint <node> [--timeout <time>] [--address <address>] [--token <hash-or-name>]",
-			Action:    getNEP17Balance,
-			Flags:     balanceFlags,
+			Description: `Prints NEP-17 balances for address and tokens specified. By default (no
+   address or token parameter) all tokens for all accounts in the specified wallet
+   are listed. A single account can be chosen with the address option and/or a
+   single token can be selected with the token option. Tokens can be specified
+   by hash, address, name or symbol. Hashes and addresses always work (as long
+   as they belong to a correct NEP-17 contract), while names or symbols (if
+   they're not NEO or GAS names/symbols) are matched against the token data
+   stored in the wallet (see import command) or balance data returned from the
+   server. If the token is not specified directly (with hash/address) and is
+   not found in the wallet then depending on the balances data from the server
+   this command can print no data at all or print multiple tokens for one
+   account (if they use the same names/symbols).
+`,
+			Action: getNEP17Balance,
+			Flags:  balanceFlags,
 		},
 		{
 			Name:      "import",
@@ -192,12 +205,42 @@ func getNEP17Balance(ctx *cli.Context) error {
 	}
 
 	name := ctx.String("token")
+	var token *wallet.Token
+
+	if name != "" {
+		// Token was explicitly specified, let's try finding it, search in the wallet first.
+		token, err = getMatchingToken(ctx, wall, name, manifest.NEP17StandardName)
+		if err != nil {
+			var h util.Uint160
+
+			// Well-known hardcoded names/symbols.
+			if name == nativenames.Neo || name == "NEO" {
+				h = neo.Hash
+			} else if name == nativenames.Gas || name == "GAS" {
+				h = gas.Hash
+			} else {
+				// The last resort, maybe it's a direct hash or address.
+				h, _ = flags.ParseAddress(name)
+			}
+			// If the hash is not found then it's some kind of named token, there is
+			// no way for us to find it, but it's not an error, maybe we'll find it
+			// in balances.
+			if !h.Equals(util.Uint160{}) {
+				// But if we have an exact hash, it must be correct.
+				token, err = getTokenWithStandard(c, h, manifest.NEP17StandardName)
+				if err != nil {
+					return cli.NewExitError(fmt.Errorf("%q is not a valid NEP-17 token: %w", name, err), 1)
+				}
+			}
+		}
+	}
 
 	for k, acc := range accounts {
 		addrHash, err := address.StringToUint160(acc.Address)
 		if err != nil {
 			return cli.NewExitError(fmt.Errorf("invalid account address: %w", err), 1)
 		}
+		// We can't use nep17.BalanceOf() even if the token is known because of LastUpdated.
 		balances, err := c.GetNEP17Balances(addrHash)
 		if err != nil {
 			return cli.NewExitError(err, 1)
@@ -210,8 +253,12 @@ func getNEP17Balance(ctx *cli.Context) error {
 
 		var tokenFound bool
 		for i := range balances.Balances {
-			token := tokenFromNEP17Balance(&balances.Balances[i])
-			if name != "" && !(token.Name == name || token.Symbol == name || token.Address() == name || token.Hash.StringLE() == name) {
+			curToken := tokenFromNEP17Balance(&balances.Balances[i])
+			if token != nil && token.Hash != curToken.Hash {
+				continue
+			}
+			// Loose (named non-native) token specification, try our best.
+			if name != "" && !(curToken.Name == name || curToken.Symbol == name || curToken.Address() == name || curToken.Hash.StringLE() == name) {
 				continue
 			}
 			printAssetBalance(ctx, balances.Balances[i])
@@ -220,52 +267,20 @@ func getNEP17Balance(ctx *cli.Context) error {
 		if name == "" || tokenFound {
 			continue
 		}
-		// Token was explicitly specified, but was not found among balances, thus either balance is 0
-		// or the token doesn't exist. Try to find token by its address/hash/name/symbol and print zero
-		// balance if found. Search into wallet first.
-		token, err := getMatchingToken(ctx, wall, name, manifest.NEP17StandardName)
-		if err != nil {
-			// The wallet doesn't contain specified token, so try to ask chain.
-			h, err := flags.ParseAddress(name)
-			if err != nil {
-				h, err = c.GetNativeContractHash(name)
-				if err != nil {
-					// Try to get native NEP17 with matching symbol.
-					var gasSymbol, neoSymbol string
-					g := gas.NewReader(invoker.New(c, nil))
-					gasSymbol, err = g.Symbol()
-					if err != nil {
-						continue
-					}
-					if gasSymbol != name {
-						n := neo.NewReader(invoker.New(c, nil))
-						neoSymbol, err = n.Symbol()
-						if err != nil {
-							continue
-						}
-						if neoSymbol != name {
-							continue
-						} else {
-							h = neo.Hash
-						}
-					} else {
-						h = gas.Hash
-					}
-				}
-			}
-			token, err = getTokenWithStandard(c, h, manifest.NEP17StandardName)
-			if err != nil {
-				continue
-			}
+		if token != nil {
+			// We have an exact token, but there is no balance data for it -> print 0.
+			printAssetBalance(ctx, result.NEP17Balance{
+				Asset:       token.Hash,
+				Amount:      "0",
+				Decimals:    int(token.Decimals),
+				LastUpdated: 0,
+				Name:        token.Name,
+				Symbol:      token.Symbol,
+			})
+		} else {
+			// We have no data for this token at all, maybe it's not even correct -> complain.
+			fmt.Fprintf(ctx.App.Writer, "Can't find data for %q token\n", name)
 		}
-		printAssetBalance(ctx, result.NEP17Balance{
-			Asset:       token.Hash,
-			Amount:      "0",
-			Decimals:    int(token.Decimals),
-			LastUpdated: 0,
-			Name:        token.Name,
-			Symbol:      token.Symbol,
-		})
 	}
 	return nil
 }
diff --git a/docs/cli.md b/docs/cli.md
index 769597d76..fdfb4651b 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -403,9 +403,10 @@ Getting balance is easy:
 ./bin/neo-go wallet nep17 balance -w /etc/neo-go/wallet.json -r http://localhost:20332
 ```
 
-By default, you'll get data for all tokens for the default wallet's
-address. You can select non-default address with `-a` flag and/or select token
-with `--token` flag (token hash or name can be used as parameter).
+By default, you'll get data for all tokens that are owned by all accounts
+stored in the given wallet. You can specify a particular address with `-a`
+flag and/or select token with `--token` flag (token hash, address, name or
+symbol can be used as a parameter).
 
 #### Transfers