From cf3ec6d9ac0621ab3ed60ecae4df8dc253b71af9 Mon Sep 17 00:00:00 2001
From: AnnaShaleva <shaleva.ann@gmail.com>
Date: Mon, 27 Dec 2021 16:29:35 +0300
Subject: [PATCH] cli: print zero balance of known token if `token` flag
 specified

Close #2313.
---
 cli/nep17_test.go   | 10 ++++++
 cli/wallet/nep17.go | 76 ++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 78 insertions(+), 8 deletions(-)

diff --git a/cli/nep17_test.go b/cli/nep17_test.go
index 605470775..32ef5be5f 100644
--- a/cli/nep17_test.go
+++ b/cli/nep17_test.go
@@ -49,6 +49,16 @@ func TestNEP17Balance(t *testing.T) {
 		b := e.Chain.GetUtilityTokenBalance(validatorHash)
 		e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+fixedn.Fixed8(b.Int64()).String()+"$")
 	})
+	t.Run("zero balance of known token", func(t *testing.T) {
+		e.Run(t, append(cmdbase, []string{"--token", "NEO"}...)...)
+		addr1, err := address.StringToUint160("Nhfg3TbpwogLvDGVvAvqyThbsHgoSUKwtn")
+		require.NoError(t, err)
+		e.checkNextLine(t, "^Account "+address.Uint160ToString(addr1))
+		e.checkNextLine(t, "^\\s*NEO:\\s+NeoToken \\("+e.Chain.GoverningTokenHash().StringLE()+"\\)")
+		e.checkNextLine(t, "^\\s*Amount\\s*:\\s*"+fixedn.Fixed8(0).String()+"$")
+		e.checkNextLine(t, "^\\s*Updated:")
+		e.checkNextLine(t, "^\\s*$")
+	})
 	t.Run("all accounts", func(t *testing.T) {
 		e.Run(t, cmdbase...)
 		addr1, err := address.StringToUint160("Nhfg3TbpwogLvDGVvAvqyThbsHgoSUKwtn")
diff --git a/cli/wallet/nep17.go b/cli/wallet/nep17.go
index 2fec82629..ecb055b3c 100644
--- a/cli/wallet/nep17.go
+++ b/cli/wallet/nep17.go
@@ -11,9 +11,11 @@ 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/encoding/address"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/fixedn"
 	"github.com/nspcc-dev/neo-go/pkg/rpc/client"
+	"github.com/nspcc-dev/neo-go/pkg/rpc/response/result"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/wallet"
@@ -188,6 +190,7 @@ func getNEP17Balance(ctx *cli.Context) error {
 		}
 		fmt.Fprintf(ctx.App.Writer, "Account %s\n", acc.Address)
 
+		var tokenFound bool
 		for i := range balances.Balances {
 			var tokenName, tokenSymbol string
 			tokenDecimals := 0
@@ -203,27 +206,84 @@ func getNEP17Balance(ctx *cli.Context) error {
 				tokenName = token.Name
 				tokenSymbol = token.Symbol
 				tokenDecimals = int(token.Decimals)
+				tokenFound = true
 			} else {
 				if name != "" {
 					continue
 				}
 				tokenSymbol = "UNKNOWN"
 			}
-			fmt.Fprintf(ctx.App.Writer, "%s: %s (%s)\n", tokenSymbol, tokenName, asset.StringLE())
-			amount := balances.Balances[i].Amount
-			if tokenDecimals != 0 {
-				b, ok := new(big.Int).SetString(amount, 10)
-				if ok {
-					amount = fixedn.ToString(b, tokenDecimals)
+			printAssetBalance(ctx, asset, tokenName, tokenSymbol, tokenDecimals, balances.Balances[i])
+		}
+		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
+					gasSymbol, h, err = getNativeNEP17Symbol(c, nativenames.Gas)
+					if err != nil {
+						continue
+					}
+					if gasSymbol != name {
+						neoSymbol, h, err = getNativeNEP17Symbol(c, nativenames.Neo)
+						if err != nil {
+							continue
+						}
+						if neoSymbol != name {
+							continue
+						}
+					}
 				}
 			}
-			fmt.Fprintf(ctx.App.Writer, "\tAmount : %s\n", amount)
-			fmt.Fprintf(ctx.App.Writer, "\tUpdated: %d\n", balances.Balances[i].LastUpdated)
+			token, err = c.NEP17TokenInfo(h)
+			if err != nil {
+				continue
+			}
 		}
+		printAssetBalance(ctx, token.Hash, token.Name, token.Symbol, int(token.Decimals), result.NEP17Balance{
+			Asset:       token.Hash,
+			Amount:      "0",
+			LastUpdated: 0,
+		})
 	}
 	return nil
 }
 
+func printAssetBalance(ctx *cli.Context, asset util.Uint160, tokenName, tokenSymbol string, tokenDecimals int, balance result.NEP17Balance) {
+	fmt.Fprintf(ctx.App.Writer, "%s: %s (%s)\n", tokenSymbol, tokenName, asset.StringLE())
+	amount := balance.Amount
+	if tokenDecimals != 0 {
+		b, ok := new(big.Int).SetString(amount, 10)
+		if ok {
+			amount = fixedn.ToString(b, tokenDecimals)
+		}
+	}
+	fmt.Fprintf(ctx.App.Writer, "\tAmount : %s\n", amount)
+	fmt.Fprintf(ctx.App.Writer, "\tUpdated: %d\n", balance.LastUpdated)
+}
+
+func getNativeNEP17Symbol(c *client.Client, name string) (string, util.Uint160, error) {
+	h, err := c.GetNativeContractHash(name)
+	if err != nil {
+		return "", util.Uint160{}, fmt.Errorf("failed to get native %s hash: %w", name, err)
+	}
+	symbol, err := c.NEP17Symbol(h)
+	if err != nil {
+		return "", util.Uint160{}, fmt.Errorf("failed to get native %s symbol: %w", name, err)
+	}
+	return symbol, h, nil
+}
+
 func getMatchingToken(ctx *cli.Context, w *wallet.Wallet, name string, standard string) (*wallet.Token, error) {
 	return getMatchingTokenAux(ctx, func(i int) *wallet.Token {
 		return w.Extra.Tokens[i]