From a61a3d5ceb18e304d8bd5dd7198f936f09d2ca0f Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Thu, 22 Apr 2021 18:21:18 +0300
Subject: [PATCH] cli: add `nep11 import` command

---
 cli/nep11_test.go    | 41 +++++++++++++++++++++++++++++++++++++++++
 cli/nep17_test.go    |  8 ++++++++
 cli/wallet/nep11.go  | 22 ++++++++++++++++++++++
 cli/wallet/nep17.go  | 22 +++++++++++++++++-----
 cli/wallet/wallet.go |  5 +++++
 5 files changed, 93 insertions(+), 5 deletions(-)
 create mode 100644 cli/nep11_test.go
 create mode 100644 cli/wallet/nep11.go

diff --git a/cli/nep11_test.go b/cli/nep11_test.go
new file mode 100644
index 000000000..0fd8966cd
--- /dev/null
+++ b/cli/nep11_test.go
@@ -0,0 +1,41 @@
+package main
+
+import (
+	"os"
+	"path"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/core/native/nativenames"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNEP11Import(t *testing.T) {
+	e := newExecutor(t, true)
+
+	tmpDir := os.TempDir()
+	walletPath := path.Join(tmpDir, "walletForImport.json")
+	defer os.Remove(walletPath)
+
+	nnsContractHash, err := e.Chain.GetNativeContractScriptHash(nativenames.NameService)
+	require.NoError(t, err)
+	neoContractHash, err := e.Chain.GetNativeContractScriptHash(nativenames.Neo)
+	require.NoError(t, err)
+	e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath)
+
+	args := []string{
+		"neo-go", "wallet", "nep11", "import",
+		"--rpc-endpoint", "http://" + e.RPC.Addr,
+		"--wallet", walletPath,
+	}
+	// missing token hash
+	e.RunWithError(t, args...)
+
+	// good
+	e.Run(t, append(args, "--token", nnsContractHash.StringLE())...)
+
+	// already exists
+	e.RunWithError(t, append(args, "--token", nnsContractHash.StringLE())...)
+
+	// not a NEP11 token
+	e.RunWithError(t, append(args, "--token", neoContractHash.StringLE())...)
+}
diff --git a/cli/nep17_test.go b/cli/nep17_test.go
index 5d35888a4..659343a8d 100644
--- a/cli/nep17_test.go
+++ b/cli/nep17_test.go
@@ -257,6 +257,8 @@ func TestNEP17ImportToken(t *testing.T) {
 	require.NoError(t, err)
 	gasContractHash, err := e.Chain.GetNativeContractScriptHash(nativenames.Gas)
 	require.NoError(t, err)
+	nnsContractHash, err := e.Chain.GetNativeContractScriptHash(nativenames.NameService)
+	require.NoError(t, err)
 	e.Run(t, "neo-go", "wallet", "init", "--wallet", walletPath)
 
 	// missing token hash
@@ -273,6 +275,12 @@ func TestNEP17ImportToken(t *testing.T) {
 		"--wallet", walletPath,
 		"--token", address.Uint160ToString(neoContractHash)) // try address instead of sh
 
+	// not a NEP17 token
+	e.RunWithError(t, "neo-go", "wallet", "nep17", "import",
+		"--rpc-endpoint", "http://"+e.RPC.Addr,
+		"--wallet", walletPath,
+		"--token", nnsContractHash.StringLE())
+
 	t.Run("Info", func(t *testing.T) {
 		checkGASInfo := func(t *testing.T) {
 			e.checkNextLine(t, "^Name:\\s*GasToken")
diff --git a/cli/wallet/nep11.go b/cli/wallet/nep11.go
new file mode 100644
index 000000000..81258dbcd
--- /dev/null
+++ b/cli/wallet/nep11.go
@@ -0,0 +1,22 @@
+package wallet
+
+import (
+	"github.com/nspcc-dev/neo-go/pkg/smartcontract/manifest"
+	"github.com/urfave/cli"
+)
+
+func newNEP11Commands() []cli.Command {
+	return []cli.Command{
+		{
+			Name:      "import",
+			Usage:     "import NEP11 token to a wallet",
+			UsageText: "import --wallet <path> --rpc-endpoint <node> --timeout <time> --token <hash>",
+			Action:    importNEP11Token,
+			Flags:     importFlags,
+		},
+	}
+}
+
+func importNEP11Token(ctx *cli.Context) error {
+	return importNEPToken(ctx, manifest.NEP11StandardName)
+}
diff --git a/cli/wallet/nep17.go b/cli/wallet/nep17.go
index 97175b5a6..941670de0 100644
--- a/cli/wallet/nep17.go
+++ b/cli/wallet/nep17.go
@@ -242,19 +242,23 @@ func getMatchingTokenAux(ctx *cli.Context, get func(i int) *wallet.Token, n int,
 			if count == 1 {
 				printTokenInfo(ctx, token)
 				printTokenInfo(ctx, t)
-				return nil, errors.New("multiple matching tokens found")
+				return nil, fmt.Errorf("multiple matching %s tokens found", standard)
 			}
 			count++
 			token = t
 		}
 	}
 	if count == 0 {
-		return nil, errors.New("token was not found")
+		return nil, fmt.Errorf("%s token was not found", standard)
 	}
 	return token, nil
 }
 
 func importNEP17Token(ctx *cli.Context) error {
+	return importNEPToken(ctx, manifest.NEP17StandardName)
+}
+
+func importNEPToken(ctx *cli.Context, standard string) error {
 	wall, err := openWallet(ctx.String("wallet"))
 	if err != nil {
 		return cli.NewExitError(err, 1)
@@ -268,9 +272,9 @@ func importNEP17Token(ctx *cli.Context) error {
 	tokenHash := tokenHashFlag.Uint160()
 
 	for _, t := range wall.Extra.Tokens {
-		if t.Hash.Equals(tokenHash) {
+		if t.Hash.Equals(tokenHash) && t.Standard == standard {
 			printTokenInfo(ctx, t)
-			return cli.NewExitError("token already exists", 1)
+			return cli.NewExitError(fmt.Errorf("%s token already exists", standard), 1)
 		}
 	}
 
@@ -282,7 +286,15 @@ func importNEP17Token(ctx *cli.Context) error {
 		return cli.NewExitError(err, 1)
 	}
 
-	tok, err := c.NEP17TokenInfo(tokenHash)
+	var tok *wallet.Token
+	switch standard {
+	case manifest.NEP17StandardName:
+		tok, err = c.NEP17TokenInfo(tokenHash)
+	case manifest.NEP11StandardName:
+		tok, err = c.NEP11TokenInfo(tokenHash)
+	default:
+		return cli.NewExitError(fmt.Sprintf("unsupported token standard: %s", standard), 1)
+	}
 	if err != nil {
 		return cli.NewExitError(fmt.Errorf("can't receive token info: %w", err), 1)
 	}
diff --git a/cli/wallet/wallet.go b/cli/wallet/wallet.go
index 9b2f3eee7..82e165925 100644
--- a/cli/wallet/wallet.go
+++ b/cli/wallet/wallet.go
@@ -234,6 +234,11 @@ func NewCommands() []cli.Command {
 				Usage:       "work with NEP17 contracts",
 				Subcommands: newNEP17Commands(),
 			},
+			{
+				Name:        "nep11",
+				Usage:       "work with NEP11 contracts",
+				Subcommands: newNEP11Commands(),
+			},
 			{
 				Name:        "candidate",
 				Usage:       "work with candidates",