Merge pull request #685 from nspcc-dev/feature/wallet

cli: implement wallet import/export functionality, part of #26.

cli/wallet/wallet.go

@@ -8,6 +8,8 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/CityOfZion/neo-go/pkg/crypto/keys"
+	"github.com/CityOfZion/neo-go/pkg/encoding/address"
 	"github.com/CityOfZion/neo-go/pkg/wallet"
 	"github.com/urfave/cli"
 	"golang.org/x/crypto/ssh/terminal"
@@ -18,6 +20,21 @@ var (
 	errPhraseMismatch = errors.New("the entered pass-phrases do not match. Maybe you have misspelled them")
 )
 
+var (
+	walletPathFlag = cli.StringFlag{
+		Name:  "path, p",
+		Usage: "Target location of the wallet file.",
+	}
+	wifFlag = cli.StringFlag{
+		Name:  "wif",
+		Usage: "WIF to import",
+	}
+	decryptFlag = cli.BoolFlag{
+		Name:  "decrypt, d",
+		Usage: "Decrypt encrypted keys.",
+	}
+)
+
 // NewCommands returns 'wallet' command.
 func NewCommands() []cli.Command {
 	return []cli.Command{{
@@ -29,28 +46,69 @@ func NewCommands() []cli.Command {
 				Usage:  "create a new wallet",
 				Action: createWallet,
 				Flags: []cli.Flag{
-					cli.StringFlag{
-						Name:  "path, p",
-						Usage: "Target location of the wallet file.",
-					},
+					walletPathFlag,
 					cli.BoolFlag{
 						Name:  "account, a",
 						Usage: "Create a new account",
 					},
 				},
 			},
+			{
+				Name:   "create-account",
+				Usage:  "add an account to the existing wallet",
+				Action: addAccount,
+				Flags: []cli.Flag{
+					walletPathFlag,
+				},
+			},
 			{
 				Name:   "dump",
 				Usage:  "check and dump an existing NEO wallet",
 				Action: dumpWallet,
 				Flags: []cli.Flag{
+					walletPathFlag,
+					decryptFlag,
+				},
+			},
+			{
+				Name:      "export",
+				Usage:     "export keys for address",
+				UsageText: "export --path <path> [--decrypt] [<address>]",
+				Action:    exportKeys,
+				Flags: []cli.Flag{
+					walletPathFlag,
+					decryptFlag,
+				},
+			},
+			{
+				Name:   "import",
+				Usage:  "import WIF",
+				Action: importWallet,
+				Flags: []cli.Flag{
+					walletPathFlag,
+					wifFlag,
 					cli.StringFlag{
-						Name:  "path, p",
-						Usage: "Target location of the wallet file.",
+						Name:  "name, n",
+						Usage: "Optional account name",
 					},
-					cli.BoolFlag{
-						Name:  "decrypt, d",
-						Usage: "Decrypt encrypted keys.",
+				},
+			},
+			{
+				Name:  "import-multisig",
+				Usage: "import multisig contract",
+				UsageText: "import-multisig --path <path> --wif <wif> --min <n>" +
+					" [<pubkey1> [<pubkey2> [...]]]",
+				Action: importMultisig,
+				Flags: []cli.Flag{
+					walletPathFlag,
+					wifFlag,
+					cli.StringFlag{
+						Name:  "name, n",
+						Usage: "Optional account name",
+					},
+					cli.IntFlag{
+						Name:  "min, m",
+						Usage: "Minimal number of signatures",
 					},
 				},
 			},
@@ -58,24 +116,152 @@ func NewCommands() []cli.Command {
 	}}
 }
 
-func dumpWallet(ctx *cli.Context) error {
-	path := ctx.String("path")
-	if len(path) == 0 {
-		return cli.NewExitError(errNoPath, 1)
+func addAccount(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
 	}
-	wall, err := wallet.NewWalletFromFile(path)
+
+	defer wall.Close()
+
+	if err := createAccount(ctx, wall); err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	return nil
+}
+
+func exportKeys(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	var addr string
+
+	decrypt := ctx.Bool("decrypt")
+	if ctx.NArg() == 0 && decrypt {
+		return cli.NewExitError(errors.New("address must be provided if '--decrypt' flag is used"), 1)
+	} else if ctx.NArg() > 0 {
+		// check address format just to catch possible typos
+		addr = ctx.Args().First()
+		_, err := address.StringToUint160(addr)
+		if err != nil {
+			return cli.NewExitError(fmt.Errorf("can't parse address: %v", err), 1)
+		}
+	}
+
+	var wifs []string
+
+loop:
+	for _, a := range wall.Accounts {
+		if addr != "" && a.Address != addr {
+			continue
+		}
+
+		for i := range wifs {
+			if a.EncryptedWIF == wifs[i] {
+				continue loop
+			}
+		}
+
+		wifs = append(wifs, a.EncryptedWIF)
+	}
+
+	for _, wif := range wifs {
+		if decrypt {
+			pass, err := readPassword("Enter password > ")
+			if err != nil {
+				return cli.NewExitError(err, 1)
+			}
+
+			pk, err := keys.NEP2Decrypt(wif, pass)
+			if err != nil {
+				return cli.NewExitError(err, 1)
+			}
+
+			wif = pk.WIF()
+		}
+
+		fmt.Println(wif)
+	}
+
+	return nil
+}
+
+func importMultisig(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	defer wall.Close()
+
+	m := ctx.Int("min")
+	if ctx.NArg() < m {
+		return cli.NewExitError(errors.New("insufficient number of public keys"), 1)
+	}
+
+	args := []string(ctx.Args())
+	pubs := make([]*keys.PublicKey, len(args))
+
+	for i := range args {
+		pubs[i], err = keys.NewPublicKeyFromString(args[i])
+		if err != nil {
+			return cli.NewExitError(fmt.Errorf("can't decode public key %d: %v", i, err), 1)
+		}
+	}
+
+	acc, err := newAccountFromWIF(ctx.String("wif"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	if err := acc.ConvertMultisig(m, pubs); err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	if err := addAccountAndSave(wall, acc); err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	return nil
+}
+
+func importWallet(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	defer wall.Close()
+
+	acc, err := newAccountFromWIF(ctx.String("wif"))
+	if err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	acc.Label = ctx.String("name")
+	if err := addAccountAndSave(wall, acc); err != nil {
+		return cli.NewExitError(err, 1)
+	}
+
+	return nil
+}
+
+func dumpWallet(ctx *cli.Context) error {
+	wall, err := openWallet(ctx.String("path"))
 	if err != nil {
 		return cli.NewExitError(err, 1)
 	}
 	if ctx.Bool("decrypt") {
-		fmt.Print("Wallet password: ")
-		pass, err := terminal.ReadPassword(int(syscall.Stdin))
+		pass, err := readPassword("Enter wallet password > ")
 		if err != nil {
 			return cli.NewExitError(err, 1)
 		}
 		for i := range wall.Accounts {
 			// Just testing the decryption here.
-			err := wall.Accounts[i].Decrypt(string(pass))
+			err := wall.Accounts[i].Decrypt(pass)
 			if err != nil {
 				return cli.NewExitError(err, 1)
 			}
@@ -109,34 +295,94 @@ func createWallet(ctx *cli.Context) error {
 	return nil
 }
 
-func createAccount(ctx *cli.Context, wall *wallet.Wallet) error {
-	var (
-		rawName,
-		rawPhrase,
-		rawPhraseCheck []byte
-	)
+func readAccountInfo() (string, string, error) {
 	buf := bufio.NewReader(os.Stdin)
 	fmt.Print("Enter the name of the account > ")
-	rawName, _ = buf.ReadBytes('\n')
-	fmt.Print("Enter passphrase > ")
-	rawPhrase, _ = terminal.ReadPassword(int(syscall.Stdin))
-	fmt.Print("\nConfirm passphrase > ")
-	rawPhraseCheck, _ = terminal.ReadPassword(int(syscall.Stdin))
-
-	// Clean data
-	var (
-		name        = strings.TrimRight(string(rawName), "\n")
-		phrase      = strings.TrimRight(string(rawPhrase), "\n")
-		phraseCheck = strings.TrimRight(string(rawPhraseCheck), "\n")
-	)
-
-	if phrase != phraseCheck {
-		return errPhraseMismatch
+	rawName, _ := buf.ReadBytes('\n')
+	phrase, err := readPassword("Enter passphrase > ")
+	if err != nil {
+		return "", "", err
+	}
+	phraseCheck, err := readPassword("Confirm passphrase > ")
+	if err != nil {
+		return "", "", err
 	}
 
+	if phrase != phraseCheck {
+		return "", "", errPhraseMismatch
+	}
+
+	name := strings.TrimRight(string(rawName), "\n")
+	return name, phrase, nil
+}
+
+func createAccount(ctx *cli.Context, wall *wallet.Wallet) error {
+	name, phrase, err := readAccountInfo()
+	if err != nil {
+		return err
+	}
 	return wall.CreateAccount(name, phrase)
 }
 
+func openWallet(path string) (*wallet.Wallet, error) {
+	if len(path) == 0 {
+		return nil, errNoPath
+	}
+	return wallet.NewWalletFromFile(path)
+}
+
+func newAccountFromWIF(wif string) (*wallet.Account, error) {
+	// note: NEP2 strings always have length of 58 even though
+	// base58 strings can have different lengths even if slice lengths are equal
+	if len(wif) == 58 {
+		pass, err := readPassword("Enter password > ")
+		if err != nil {
+			return nil, err
+		}
+
+		return wallet.NewAccountFromEncryptedWIF(wif, pass)
+	}
+
+	acc, err := wallet.NewAccountFromWIF(wif)
+	if err != nil {
+		return nil, err
+	}
+
+	fmt.Println("Provided WIF was unencrypted. Wallet can contain only encrypted keys.")
+	name, pass, err := readAccountInfo()
+	if err != nil {
+		return nil, err
+	}
+
+	acc.Label = name
+	if err := acc.Encrypt(pass); err != nil {
+		return nil, err
+	}
+
+	return acc, nil
+}
+
+func addAccountAndSave(w *wallet.Wallet, acc *wallet.Account) error {
+	for i := range w.Accounts {
+		if w.Accounts[i].Address == acc.Address {
+			return fmt.Errorf("address '%s' is already in wallet", acc.Address)
+		}
+	}
+
+	w.AddAccount(acc)
+	return w.Save()
+}
+
+func readPassword(prompt string) (string, error) {
+	fmt.Print(prompt)
+	rawPass, err := terminal.ReadPassword(syscall.Stdin)
+	fmt.Println()
+	if err != nil {
+		return "", err
+	}
+	return strings.TrimRight(string(rawPass), "\n"), nil
+}
+
 func fmtPrintWallet(wall *wallet.Wallet) {
 	b, _ := wall.JSON()
 	fmt.Println("")

pkg/smartcontract/param_context.go

@@ -2,6 +2,7 @@ package smartcontract
 
 import (
 	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"strconv"
 	"strings"
@@ -78,6 +79,22 @@ func (pt ParamType) MarshalJSON() ([]byte, error) {
 	return []byte(`"` + pt.String() + `"`), nil
 }
 
+// UnmarshalJSON implements json.Unmarshaler interface.
+func (pt *ParamType) UnmarshalJSON(data []byte) error {
+	var s string
+	if err := json.Unmarshal(data, &s); err != nil {
+		return err
+	}
+
+	p, err := parseParamType(s)
+	if err != nil {
+		return err
+	}
+
+	*pt = p
+	return nil
+}
+
 // EncodeBinary implements io.Serializable interface.
 func (pt ParamType) EncodeBinary(w *io.BinWriter) {
 	w.WriteB(byte(pt))

pkg/wallet/account.go

@@ -1,12 +1,16 @@
 package wallet
 
 import (
+	"bytes"
 	"encoding/hex"
 	"encoding/json"
 	"errors"
+	"fmt"
 
 	"github.com/CityOfZion/neo-go/pkg/crypto/hash"
 	"github.com/CityOfZion/neo-go/pkg/crypto/keys"
+	"github.com/CityOfZion/neo-go/pkg/encoding/address"
+	"github.com/CityOfZion/neo-go/pkg/smartcontract"
 	"github.com/CityOfZion/neo-go/pkg/util"
 )
 
@@ -50,7 +54,7 @@ type Contract struct {
 	Script []byte `json:"script"`
 
 	// A list of parameters used deploying this contract.
-	Parameters []interface{} `json:"parameters"`
+	Parameters []contractParam `json:"parameters"`
 
 	// Indicates whether the contract has been deployed to the blockchain.
 	Deployed bool `json:"deployed"`
@@ -62,12 +66,17 @@ type contract struct {
 	Script string `json:"script"`
 
 	// A list of parameters used deploying this contract.
-	Parameters []interface{} `json:"parameters"`
+	Parameters []contractParam `json:"parameters"`
 
 	// Indicates whether the contract has been deployed to the blockchain.
 	Deployed bool `json:"deployed"`
 }
 
+type contractParam struct {
+	Name string                  `json:"name"`
+	Type smartcontract.ParamType `json:"type"`
+}
+
 // ScriptHash returns the hash of contract's script.
 func (c Contract) ScriptHash() util.Uint160 {
 	return hash.Hash160(c.Script)
@@ -122,7 +131,14 @@ func (a *Account) Decrypt(passphrase string) error {
 		return errors.New("no encrypted wif in the account")
 	}
 	a.privateKey, err = keys.NEP2Decrypt(a.EncryptedWIF, passphrase)
-	return err
+	if err != nil {
+		return err
+	}
+
+	a.publicKey = a.privateKey.PublicKey().Bytes()
+	a.wif = a.privateKey.WIF()
+
+	return nil
 }
 
 // Encrypt encrypts the wallet's PrivateKey with the given passphrase
@@ -150,6 +166,47 @@ func NewAccountFromWIF(wif string) (*Account, error) {
 	return newAccountFromPrivateKey(privKey), nil
 }
 
+// NewAccountFromEncryptedWIF creates a new Account from the given encrypted WIF.
+func NewAccountFromEncryptedWIF(wif string, pass string) (*Account, error) {
+	priv, err := keys.NEP2Decrypt(wif, pass)
+	if err != nil {
+		return nil, err
+	}
+
+	a := newAccountFromPrivateKey(priv)
+	a.EncryptedWIF = wif
+
+	return a, nil
+}
+
+// ConvertMultisig sets a's contract to multisig contract with m sufficient signatures.
+func (a *Account) ConvertMultisig(m int, pubs []*keys.PublicKey) error {
+	var found bool
+	for i := range pubs {
+		if bytes.Equal(a.publicKey, pubs[i].Bytes()) {
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		return errors.New("own public key was not found among multisig keys")
+	}
+
+	script, err := smartcontract.CreateMultiSigRedeemScript(m, pubs)
+	if err != nil {
+		return err
+	}
+
+	a.Address = address.Uint160ToString(hash.Hash160(script))
+	a.Contract = &Contract{
+		Script:     script,
+		Parameters: getContractParams(m),
+	}
+
+	return nil
+}
+
 // newAccountFromPrivateKey creates a wallet from the given PrivateKey.
 func newAccountFromPrivateKey(p *keys.PrivateKey) *Account {
 	pubKey := p.PublicKey()
@@ -161,7 +218,21 @@ func newAccountFromPrivateKey(p *keys.PrivateKey) *Account {
 		privateKey: p,
 		Address:    pubAddr,
 		wif:        wif,
+		Contract: &Contract{
+			Script:     pubKey.GetVerificationScript(),
+			Parameters: getContractParams(1),
+		},
 	}
 
 	return a
 }
+
+func getContractParams(n int) []contractParam {
+	params := make([]contractParam, n)
+	for i := range params {
+		params[i].Name = fmt.Sprintf("parameter%d", i)
+		params[i].Type = smartcontract.SignatureType
+	}
+
+	return params
+}

pkg/wallet/account_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/CityOfZion/neo-go/pkg/crypto/hash"
+	"github.com/CityOfZion/neo-go/pkg/crypto/keys"
 	"github.com/CityOfZion/neo-go/pkg/internal/keytestcases"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -49,10 +50,23 @@ func TestNewFromWif(t *testing.T) {
 	}
 }
 
+func TestNewAccountFromEncryptedWIF(t *testing.T) {
+	for _, tc := range keytestcases.Arr {
+		acc, err := NewAccountFromEncryptedWIF(tc.EncryptedWif, tc.Passphrase)
+		if tc.Invalid {
+			assert.Error(t, err)
+			continue
+		}
+
+		assert.NoError(t, err)
+		compareFields(t, tc, acc)
+	}
+}
+
 func TestContract_MarshalJSON(t *testing.T) {
 	var c Contract
 
-	data := []byte(`{"script":"0102","parameters":[1],"deployed":false}`)
+	data := []byte(`{"script":"0102","parameters":[{"name":"name0", "type":"Signature"}],"deployed":false}`)
 	require.NoError(t, json.Unmarshal(data, &c))
 	require.Equal(t, []byte{1, 2}, c.Script)
 
@@ -74,6 +88,51 @@ func TestContract_ScriptHash(t *testing.T) {
 	require.Equal(t, hash.Hash160(script), c.ScriptHash())
 }
 
+func TestAccount_ConvertMultisig(t *testing.T) {
+	// test is based on a wallet1_solo.json accounts from neo-local
+	a, err := NewAccountFromEncryptedWIF("6PYLmjBYJ4wQTCEfqvnznGJwZeW9pfUcV5m5oreHxqryUgqKpTRAFt9L8Y", "one")
+	require.NoError(t, err)
+
+	hexs := []string{
+		"02b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc2", // <- this is our key
+		"02103a7f7dd016558597f7960d27c516a4394fd968b9e65155eb4b013e4040406e",
+		"02a7bc55fe8684e0119768d104ba30795bdcc86619e864add26156723ed185cd62",
+		"03d90c07df63e690ce77912e10ab51acc944b66860237b608c4f8f8309e71ee699",
+	}
+
+	t.Run("invalid number of signatures", func(t *testing.T) {
+		pubs := convertPubs(t, hexs)
+		require.Error(t, a.ConvertMultisig(0, pubs))
+	})
+
+	t.Run("account key is missing from multisig", func(t *testing.T) {
+		pubs := convertPubs(t, hexs[1:])
+		require.Error(t, a.ConvertMultisig(1, pubs))
+	})
+
+	t.Run("1/1 multisig", func(t *testing.T) {
+		pubs := convertPubs(t, hexs[:1])
+		require.NoError(t, a.ConvertMultisig(1, pubs))
+		require.Equal(t, "AbU69m8WUZJSWanfr1Cy66cpEcsmMcX7BR", a.Address)
+	})
+
+	t.Run("3/4 multisig", func(t *testing.T) {
+		pubs := convertPubs(t, hexs)
+		require.NoError(t, a.ConvertMultisig(3, pubs))
+		require.Equal(t, "AZ81H31DMWzbSnFDLFkzh9vHwaDLayV7fU", a.Address)
+	})
+}
+
+func convertPubs(t *testing.T, hexKeys []string) []*keys.PublicKey {
+	pubs := make([]*keys.PublicKey, len(hexKeys))
+	for i := range pubs {
+		var err error
+		pubs[i], err = keys.NewPublicKeyFromString(hexKeys[i])
+		require.NoError(t, err)
+	}
+	return pubs
+}
+
 func compareFields(t *testing.T, tk keytestcases.Ktype, acc *Account) {
 	if want, have := tk.Address, acc.Address; want != have {
 		t.Fatalf("expected %s got %s", want, have)

pkg/wallet/wallet.go

@@ -105,6 +105,11 @@ func (w *Wallet) Path() string {
 // that is responsible for saving the data. This can
 // be a buffer, file, etc..
 func (w *Wallet) Save() error {
+	if s, ok := w.rw.(io.Seeker); ok {
+		if _, err := s.Seek(0, 0); err != nil {
+			return err
+		}
+	}
 	return json.NewEncoder(w.rw).Encode(w)
 }
 

pkg/wallet/wallet_test.go

@@ -90,6 +90,17 @@ func TestSave(t *testing.T) {
 	openedWallet, err := NewWalletFromFile(wallet.path)
 	require.NoError(t, err)
 	require.Equal(t, wallet.Accounts, openedWallet.Accounts)
+
+	t.Run("change and rewrite", func(t *testing.T) {
+		err := openedWallet.CreateAccount("test", "pass")
+		require.NoError(t, err)
+
+		w2, err := NewWalletFromFile(openedWallet.path)
+		require.NoError(t, err)
+		require.Equal(t, 2, len(w2.Accounts))
+		require.NoError(t, w2.Accounts[1].Decrypt("pass"))
+		require.Equal(t, openedWallet.Accounts, w2.Accounts)
+	})
 }
 
 func TestJSONMarshallUnmarshal(t *testing.T) {