Merge pull request #728 from nspcc-dev/feature/nep5cli

cli: implement NEP5-related commands

cli/wallet/nep5.go

@@ -0,0 +1,344 @@
+package wallet
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/nspcc-dev/neo-go/cli/flags"
+	"github.com/nspcc-dev/neo-go/pkg/core"
+	"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/io"
+	"github.com/nspcc-dev/neo-go/pkg/rpc/client"
+	"github.com/nspcc-dev/neo-go/pkg/rpc/request"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
+	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
+	"github.com/nspcc-dev/neo-go/pkg/wallet"
+	"github.com/urfave/cli"
+)
+
+func newNEP5Commands() []cli.Command {
+	return []cli.Command{
+		{
+			Name:      "balance",
+			Usage:     "get address balance",
+			UsageText: "balance --path <path> --rpc <node> --addr <addr> [--token <hash-or-name>]",
+			Action:    getNEP5Balance,
+			Flags: []cli.Flag{
+				walletPathFlag,
+				rpcFlag,
+				timeoutFlag,
+				cli.StringFlag{
+					Name:  "addr",
+					Usage: "Address to use",
+				},
+				cli.StringFlag{
+					Name:  "token",
+					Usage: "Token to use",
+				},
+			},
+		},
+		{
+			Name:      "import",
+			Usage:     "import NEP5 token to a wallet",
+			UsageText: "import --path <path> --rpc <node> --token <hash>",
+			Action:    importNEP5Token,
+			Flags: []cli.Flag{
+				walletPathFlag,
+				rpcFlag,
+				cli.StringFlag{
+					Name:  "token",
+					Usage: "Token contract hash in LE",
+				},
+			},
+		},
+		{
+			Name:      "info",
+			Usage:     "print imported NEP5 token info",
+			UsageText: "print --path <path> [--token <hash-or-name>]",
+			Action:    printNEP5Info,
+			Flags: []cli.Flag{
+				walletPathFlag,
+				cli.StringFlag{
+					Name:  "token",
+					Usage: "Token name or hash",
+				},
+			},
+		},
+		{
+			Name:      "transfer",
+			Usage:     "transfer NEP5 tokens",
+			UsageText: "transfer --path <path> --rpc <node> --from <addr> --to <addr> --token <hash> --amount string",
+			Action:    transferNEP5,
+			Flags: []cli.Flag{
+				walletPathFlag,
+				rpcFlag,
+				timeoutFlag,
+				fromAddrFlag,
+				toAddrFlag,
+				cli.StringFlag{
+					Name:  "token",
+					Usage: "Token to use",
+				},
+				cli.StringFlag{
+					Name:  "amount",
+					Usage: "Amount of asset to send",
+				},
+				cli.StringFlag{
+					Name:  "gas",
+					Usage: "Amount of GAS to attach to a tx",
+				},
+			},
+		},
+	}
+}
+
+func getNEP5Balance(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	defer wall.Close()
+
+	addr := ctx.String("addr")
+	addrHash, err := address.StringToUint160(addr)
+	if err != nil {
+		return cli.NewExitError(fmt.Errorf("invalid address: %v", err), 1)
+	}
+	acc := wall.GetAccount(addrHash)
+	if acc == nil {
+		return cli.NewExitError(fmt.Errorf("can't find account for the address: %s", addr), 1)
+	}
+
+	gctx, cancel := getGoContext(ctx)
+	defer cancel()
+
+	c, err := client.New(gctx, ctx.String("rpc"), client.Options{})
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	var token *wallet.Token
+	name := ctx.String("token")
+	if name != "" {
+		token, err = getMatchingToken(wall, name)
+		if err != nil {
+			token, err = getMatchingTokenRPC(c, addrHash, name)
+			if err != nil {
+				return cli.NewExitError(err, 1)
+			}
+		}
+	}
+
+	balances, err := c.GetNEP5Balances(addrHash)
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	for i := range balances.Balances {
+		asset := balances.Balances[i].Asset
+		if name != "" && !token.Hash.Equals(asset) {
+			continue
+		}
+		fmt.Printf("TokenHash: %s\n", asset)
+		fmt.Printf("\tAmount : %s\n", balances.Balances[i].Amount)
+		fmt.Printf("\tUpdated: %d\n", balances.Balances[i].LastUpdated)
+	}
+	return nil
+}
+
+func getMatchingToken(w *wallet.Wallet, name string) (*wallet.Token, error) {
+	return getMatchingTokenAux(func(i int) *wallet.Token {
+		return w.Extra.Tokens[i]
+	}, len(w.Extra.Tokens), name)
+}
+
+func getMatchingTokenRPC(c *client.Client, addr util.Uint160, name string) (*wallet.Token, error) {
+	bs, err := c.GetNEP5Balances(addr)
+	if err != nil {
+		return nil, err
+	}
+	get := func(i int) *wallet.Token {
+		t, _ := c.NEP5TokenInfo(bs.Balances[i].Asset)
+		return t
+	}
+	return getMatchingTokenAux(get, len(bs.Balances), name)
+}
+
+func getMatchingTokenAux(get func(i int) *wallet.Token, n int, name string) (*wallet.Token, error) {
+	var token *wallet.Token
+	var count int
+	for i := 0; i < n; i++ {
+		t := get(i)
+		if t != nil && (t.Name == name || t.Symbol == name || t.Address == name || t.Hash.StringLE() == name) {
+			if count == 1 {
+				printTokenInfo(token)
+				printTokenInfo(t)
+				return nil, errors.New("multiple matching tokens found")
+			}
+			count++
+			token = t
+		}
+	}
+	if count == 0 {
+		return nil, errors.New("token was not found")
+	}
+	return token, nil
+}
+
+func importNEP5Token(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	defer wall.Close()
+
+	tokenHash, err := util.Uint160DecodeStringLE(ctx.String("token"))
+	if err != nil {
+		return cli.NewExitError(fmt.Errorf("invalid token contract hash: %v", err), 1)
+	}
+
+	for _, t := range wall.Extra.Tokens {
+		if t.Hash.Equals(tokenHash) {
+			printTokenInfo(t)
+			return cli.NewExitError("token already exists", 1)
+		}
+	}
+
+	gctx, cancel := getGoContext(ctx)
+	defer cancel()
+
+	c, err := client.New(gctx, ctx.String("rpc"), client.Options{})
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	tok, err := c.NEP5TokenInfo(tokenHash)
+	if err != nil {
+		return cli.NewExitError(fmt.Errorf("can't receive token info: %v", err), 1)
+	}
+
+	wall.AddToken(tok)
+	if err := wall.Save(); err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	printTokenInfo(tok)
+	return nil
+}
+
+func printTokenInfo(tok *wallet.Token) {
+	fmt.Printf("Name:\t%s\n", tok.Name)
+	fmt.Printf("Symbol:\t%s\n", tok.Symbol)
+	fmt.Printf("Hash:\t%s\n", tok.Hash.StringLE())
+	fmt.Printf("Decimals: %d\n", tok.Decimals)
+	fmt.Printf("Address: %s\n", tok.Address)
+}
+
+func printNEP5Info(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	defer wall.Close()
+
+	if name := ctx.String("token"); name != "" {
+		token, err := getMatchingToken(wall, name)
+		if err != nil {
+			return cli.NewExitError(err, 1)
+		}
+		printTokenInfo(token)
+		return nil
+	}
+
+	for i, t := range wall.Extra.Tokens {
+		if i > 0 {
+			fmt.Println()
+		}
+		printTokenInfo(t)
+	}
+	return nil
+}
+
+func transferNEP5(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+	defer wall.Close()
+
+	fromFlag := ctx.Generic("from").(*flags.Address)
+	from := fromFlag.Uint160()
+	acc := wall.GetAccount(from)
+	if acc == nil {
+		return cli.NewExitError(fmt.Errorf("can't find account for the address: %s", fromFlag), 1)
+	}
+
+	gctx, cancel := getGoContext(ctx)
+	defer cancel()
+	c, err := client.New(gctx, ctx.String("rpc"), client.Options{})
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	toFlag := ctx.Generic("to").(*flags.Address)
+	to := toFlag.Uint160()
+	token, err := getMatchingToken(wall, ctx.String("token"))
+	if err != nil {
+		fmt.Println("Can't find matching token in the wallet. Querying RPC-node for balances.")
+		token, err = getMatchingTokenRPC(c, from, ctx.String("token"))
+		if err != nil {
+			return cli.NewExitError(err, 1)
+		}
+	}
+
+	amount, err := util.FixedNFromString(ctx.String("amount"), int(token.Decimals))
+	if err != nil {
+		return cli.NewExitError(fmt.Errorf("invalid amount: %v", err), 1)
+	}
+
+	// Note: we don't use invoke function here because it requires
+	// 2 round trips instead of one.
+	w := io.NewBufBinWriter()
+	emit.Int(w.BinWriter, amount)
+	emit.Bytes(w.BinWriter, to.BytesBE())
+	emit.Bytes(w.BinWriter, from.BytesBE())
+	emit.Int(w.BinWriter, 3)
+	emit.Opcode(w.BinWriter, opcode.PACK)
+	emit.String(w.BinWriter, "transfer")
+	emit.AppCall(w.BinWriter, token.Hash, false)
+	emit.Opcode(w.BinWriter, opcode.THROWIFNOT)
+
+	var gas util.Fixed8
+	if gasString := ctx.String("gas"); gasString != "" {
+		gas, err = util.Fixed8FromString(gasString)
+		if err != nil {
+			return cli.NewExitError(fmt.Errorf("invalid GAS amount: %v", err), 1)
+		}
+	}
+
+	tx := transaction.NewInvocationTX(w.Bytes(), gas)
+	tx.Attributes = append(tx.Attributes, transaction.Attribute{
+		Usage: transaction.Script,
+		Data:  from.BytesBE(),
+	})
+
+	if err := request.AddInputsAndUnspentsToTx(tx, fromFlag.String(), core.UtilityTokenID(), gas, c); err != nil {
+		return cli.NewExitError(fmt.Errorf("can't add GAS to a tx: %v", err), 1)
+	}
+
+	if pass, err := readPassword("Password > "); err != nil {
+		return cli.NewExitError(err, 1)
+	} else if err := acc.Decrypt(pass); err != nil {
+		return cli.NewExitError(err, 1)
+	} else if err := acc.SignTx(tx); err != nil {
+		return cli.NewExitError(fmt.Errorf("can't sign tx: %v", err), 1)
+	}
+
+	if err := c.SendRawTransaction(tx); err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	fmt.Println(tx.Hash())
+	return nil
+}

cli/wallet/wallet.go

@@ -59,6 +59,14 @@ var (
 		Name:  "in",
 		Usage: "file with JSON transaction",
 	}
+	fromAddrFlag = flags.AddressFlag{
+		Name:  "from",
+		Usage: "Address to send an asset from",
+	}
+	toAddrFlag = flags.AddressFlag{
+		Name:  "to",
+		Usage: "Address to send an asset to",
+	}
 )
 
 // NewCommands returns 'wallet' command.
@@ -163,14 +171,8 @@ func NewCommands() []cli.Command {
 					rpcFlag,
 					timeoutFlag,
 					outFlag,
-					flags.AddressFlag{
+					fromAddrFlag,
-						Name:  "from",
+					toAddrFlag,
-						Usage: "Address to send an asset from",
-					},
-					flags.AddressFlag{
-						Name:  "to",
-						Usage: "Address to send an asset to",
-					},
 					cli.StringFlag{
 						Name:  "amount",
 						Usage: "Amount of asset to send",
@@ -186,6 +188,11 @@ func NewCommands() []cli.Command {
 				Usage:       "work with multisig address",
 				Subcommands: newMultisigCommands(),
 			},
+			{
+				Name:        "nep5",
+				Usage:       "work with NEP5 contracts",
+				Subcommands: newNEP5Commands(),
+			},
 		},
 	}}
 }
@@ -221,7 +228,7 @@ func claimGas(ctx *cli.Context) error {
 	if err != nil {
 		return cli.NewExitError(err, 1)
 	}
-	info, err := c.GetClaimable(scriptHash.String())
+	info, err := c.GetClaimable(addrFlag.String())
 	if err != nil {
 		return cli.NewExitError(err, 1)
 	} else if info.Unclaimed == 0 || len(info.Spents) == 0 {

pkg/core/state/nep5_test.go

@@ -58,6 +58,12 @@ func TestNEP5Transfer_DecodeBinary(t *testing.T) {
 	testEncodeDecode(t, expected, new(NEP5Transfer))
 }
 
+func TestNEP5TransferSize(t *testing.T) {
+	tr := randomTransfer(t, rand.New(rand.NewSource(0)))
+	size := io.GetVarSize(tr)
+	require.EqualValues(t, NEP5TransferSize, size)
+}
+
 func randomTransfer(t *testing.T, r *rand.Rand) *NEP5Transfer {
 	tr := &NEP5Transfer{
 		Amount: int64(r.Uint64()),

pkg/rpc/client/nep5.go

@@ -0,0 +1,132 @@
+package client
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
+	"github.com/nspcc-dev/neo-go/pkg/wallet"
+)
+
+// NEP5Decimals invokes `decimals` NEP5 method on a specified contract.
+func (c *Client) NEP5Decimals(tokenHash util.Uint160) (int64, error) {
+	result, err := c.InvokeFunction(tokenHash.StringLE(), "decimals", []smartcontract.Parameter{})
+	if err != nil {
+		return 0, err
+	} else if result.State != "HALT" || len(result.Stack) == 0 {
+		return 0, errors.New("invalid VM state")
+	}
+
+	return topIntFromStack(result.Stack)
+}
+
+// NEP5Name invokes `name` NEP5 method on a specified contract.
+func (c *Client) NEP5Name(tokenHash util.Uint160) (string, error) {
+	result, err := c.InvokeFunction(tokenHash.StringLE(), "name", []smartcontract.Parameter{})
+	if err != nil {
+		return "", err
+	} else if result.State != "HALT" || len(result.Stack) == 0 {
+		return "", errors.New("invalid VM state")
+	}
+
+	return topStringFromStack(result.Stack)
+}
+
+// NEP5Symbol invokes `symbol` NEP5 method on a specified contract.
+func (c *Client) NEP5Symbol(tokenHash util.Uint160) (string, error) {
+	result, err := c.InvokeFunction(tokenHash.StringLE(), "symbol", []smartcontract.Parameter{})
+	if err != nil {
+		return "", err
+	} else if result.State != "HALT" || len(result.Stack) == 0 {
+		return "", errors.New("invalid VM state")
+	}
+
+	return topStringFromStack(result.Stack)
+}
+
+// NEP5TotalSupply invokes `totalSupply` NEP5 method on a specified contract.
+func (c *Client) NEP5TotalSupply(tokenHash util.Uint160) (int64, error) {
+	result, err := c.InvokeFunction(tokenHash.StringLE(), "totalSupply", []smartcontract.Parameter{})
+	if err != nil {
+		return 0, err
+	} else if result.State != "HALT" || len(result.Stack) == 0 {
+		return 0, errors.New("invalid VM state")
+	}
+
+	return topIntFromStack(result.Stack)
+}
+
+// NEP5BalanceOf invokes `balanceOf` NEP5 method on a specified contract.
+func (c *Client) NEP5BalanceOf(tokenHash util.Uint160) (int64, error) {
+	result, err := c.InvokeFunction(tokenHash.StringLE(), "balanceOf", []smartcontract.Parameter{})
+	if err != nil {
+		return 0, err
+	} else if result.State != "HALT" || len(result.Stack) == 0 {
+		return 0, errors.New("invalid VM state")
+	}
+
+	return topIntFromStack(result.Stack)
+}
+
+// NEP5TokenInfo returns full NEP5 token info.
+func (c *Client) NEP5TokenInfo(tokenHash util.Uint160) (*wallet.Token, error) {
+	name, err := c.NEP5Name(tokenHash)
+	if err != nil {
+		return nil, err
+	}
+	symbol, err := c.NEP5Symbol(tokenHash)
+	if err != nil {
+		return nil, err
+	}
+	decimals, err := c.NEP5Decimals(tokenHash)
+	if err != nil {
+		return nil, err
+	}
+	return wallet.NewToken(tokenHash, name, symbol, decimals), nil
+}
+
+func topIntFromStack(st []smartcontract.Parameter) (int64, error) {
+	index := len(st) - 1 // top stack element is last in the array
+	var decimals int64
+	switch typ := st[index].Type; typ {
+	case smartcontract.IntegerType:
+		var ok bool
+		decimals, ok = st[index].Value.(int64)
+		if !ok {
+			return 0, errors.New("invalid Integer item")
+		}
+	case smartcontract.ByteArrayType:
+		data, ok := st[index].Value.([]byte)
+		if !ok {
+			return 0, errors.New("invalid ByteArray item")
+		}
+		decimals = emit.BytesToInt(data).Int64()
+	default:
+		return 0, fmt.Errorf("invalid stack item type: %s", typ)
+	}
+	return decimals, nil
+}
+
+func topStringFromStack(st []smartcontract.Parameter) (string, error) {
+	index := len(st) - 1 // top stack element is last in the array
+	var s string
+	switch typ := st[index].Type; typ {
+	case smartcontract.StringType:
+		var ok bool
+		s, ok = st[index].Value.(string)
+		if !ok {
+			return "", errors.New("invalid String item")
+		}
+	case smartcontract.ByteArrayType:
+		data, ok := st[index].Value.([]byte)
+		if !ok {
+			return "", errors.New("invalid ByteArray item")
+		}
+		s = string(data)
+	default:
+		return "", fmt.Errorf("invalid stack item type: %s", typ)
+	}
+	return s, nil
+}

pkg/util/fixed8.go

@@ -3,6 +3,7 @@ package util
 import (
 	"encoding/json"
 	"errors"
+	"math"
 	"strconv"
 	"strings"
 
@@ -71,25 +72,36 @@ func Fixed8FromFloat(val float64) Fixed8 {
 // Fixed8FromString parses s which must be a fixed point number
 // with precision up to 10^-8
 func Fixed8FromString(s string) (Fixed8, error) {
+	num, err := FixedNFromString(s, precision)
+	if err != nil {
+		return 0, err
+	}
+	return Fixed8(num), err
+}
+
+// FixedNFromString parses s which must be a fixed point number
+// with precision 10^-d.
+func FixedNFromString(s string, precision int) (int64, error) {
 	parts := strings.SplitN(s, ".", 2)
+	d := int64(math.Pow10(precision))
 	ip, err := strconv.ParseInt(parts[0], 10, 64)
 	if err != nil {
 		return 0, errInvalidString
 	} else if len(parts) == 1 {
-		return Fixed8(ip * decimals), nil
+		return ip * d, nil
 	}
 
 	fp, err := strconv.ParseInt(parts[1], 10, 64)
-	if err != nil || fp >= decimals {
+	if err != nil || fp >= d {
 		return 0, errInvalidString
 	}
 	for i := len(parts[1]); i < precision; i++ {
 		fp *= 10
 	}
 	if ip < 0 {
-		return Fixed8(ip*decimals - fp), nil
+		return ip*d - fp, nil
 	}
-	return Fixed8(ip*decimals + fp), nil
+	return ip*d + fp, nil
 }
 
 // UnmarshalJSON implements the json unmarshaller interface.

pkg/util/fixed8_test.go

@@ -85,6 +85,20 @@ func TestFixed8FromString(t *testing.T) {
 	assert.Error(t, err)
 }
 
+func TestFixedNFromString(t *testing.T) {
+	val := "123.456"
+	num, err := FixedNFromString(val, 3)
+	require.NoError(t, err)
+	require.EqualValues(t, 123456, num)
+
+	num, err = FixedNFromString(val, 4)
+	require.NoError(t, err)
+	require.EqualValues(t, 1234560, num)
+
+	_, err = FixedNFromString(val, 2)
+	require.Error(t, err)
+}
+
 func TestSatoshi(t *testing.T) {
 	satoshif8 := Satoshi()
 	assert.Equal(t, "0.00000001", satoshif8.String())

pkg/wallet/token.go

@@ -0,0 +1,60 @@
+package wallet
+
+import (
+	"encoding/json"
+
+	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+)
+
+// Token represents imported token contract.
+type Token struct {
+	Name     string
+	Hash     util.Uint160
+	Decimals int64
+	Symbol   string
+	Address  string
+}
+
+type tokenAux struct {
+	Name     string       `json:"name"`
+	Hash     util.Uint160 `json:"script_hash"`
+	Decimals int64        `json:"decimals"`
+	Symbol   string       `json:"symbol"`
+}
+
+// NewToken returns new token contract info.
+func NewToken(tokenHash util.Uint160, name, symbol string, decimals int64) *Token {
+	return &Token{
+		Name:     name,
+		Hash:     tokenHash,
+		Decimals: decimals,
+		Symbol:   symbol,
+		Address:  address.Uint160ToString(tokenHash),
+	}
+}
+
+// MarshalJSON implements json.Marshaler interface.
+func (t *Token) MarshalJSON() ([]byte, error) {
+	m := &tokenAux{
+		Name:     t.Name,
+		Hash:     t.Hash.Reverse(), // address should be marshaled in LE but default marshaler uses BE.
+		Decimals: t.Decimals,
+		Symbol:   t.Symbol,
+	}
+	return json.Marshal(m)
+}
+
+// UnmarshalJSON implements json.Unmarshaler interface.
+func (t *Token) UnmarshalJSON(data []byte) error {
+	aux := new(tokenAux)
+	if err := json.Unmarshal(data, aux); err != nil {
+		return err
+	}
+	t.Name = aux.Name
+	t.Hash = aux.Hash.Reverse()
+	t.Decimals = aux.Decimals
+	t.Symbol = aux.Symbol
+	t.Address = address.Uint160ToString(t.Hash)
+	return nil
+}

pkg/wallet/token_test.go

@@ -0,0 +1,29 @@
+package wallet
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/stretchr/testify/require"
+)
+
+func TestToken_MarshalJSON(t *testing.T) {
+	// From the https://neo-python.readthedocs.io/en/latest/prompt.html#import-nep5-compliant-token
+	h, err := util.Uint160DecodeStringLE("f8d448b227991cf07cb96a6f9c0322437f1599b9")
+	require.NoError(t, err)
+
+	tok := NewToken(h, "NEP5 Standard", "NEP5", 8)
+	require.Equal(t, "NEP5 Standard", tok.Name)
+	require.Equal(t, "NEP5", tok.Symbol)
+	require.EqualValues(t, 8, tok.Decimals)
+	require.Equal(t, h, tok.Hash)
+	require.Equal(t, "AYhE3Svuqdfh1RtzvE8hUhNR7HSpaSDFQg", tok.Address)
+
+	data, err := json.Marshal(tok)
+	require.NoError(t, err)
+
+	actual := new(Token)
+	require.NoError(t, json.Unmarshal(data, actual))
+	require.Equal(t, tok, actual)
+}

pkg/wallet/wallet.go

@@ -28,7 +28,7 @@ type Wallet struct {
 
 	// Extra metadata can be used for storing arbitrary data.
 	// This field can be empty.
-	Extra interface{} `json:"extra"`
+	Extra Extra `json:"extra"`
 
 	// Path where the wallet file is located..
 	path string
@@ -37,6 +37,12 @@ type Wallet struct {
 	rw io.ReadWriter
 }
 
+// Extra stores imported token contracts.
+type Extra struct {
+	// Tokens is a list of imported token contracts.
+	Tokens []*Token
+}
+
 // NewWallet creates a new NEO wallet at the given location.
 func NewWallet(location string) (*Wallet, error) {
 	file, err := os.Create(location)
@@ -96,6 +102,11 @@ func (w *Wallet) AddAccount(acc *Account) {
 	w.Accounts = append(w.Accounts, acc)
 }
 
+// AddToken adds new token to a wallet.
+func (w *Wallet) AddToken(tok *Token) {
+	w.Extra.Tokens = append(w.Extra.Tokens, tok)
+}
+
 // Path returns the location of the wallet on the filesystem.
 func (w *Wallet) Path() string {
 	return w.path