github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/wallet/wallet.go (about)

     1  package wallet
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"errors"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  
    11  	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
    12  	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
    13  	"github.com/nspcc-dev/neo-go/pkg/util"
    14  	"github.com/nspcc-dev/neo-go/pkg/vm"
    15  )
    16  
    17  const (
    18  	// The current version of neo-go wallet implementations.
    19  	walletVersion = "1.0"
    20  )
    21  
    22  var (
    23  	// ErrPathIsEmpty appears if wallet was created without linking to file system path,
    24  	// for instance with [NewInMemoryWallet] or [NewWalletFromBytes].
    25  	// Despite this, there was an attempt to save it via [Wallet.Save] or [Wallet.SavePretty] without [Wallet.SetPath].
    26  	ErrPathIsEmpty = errors.New("path is empty")
    27  )
    28  
    29  // Wallet represents a NEO (NEP-2, NEP-6) compliant wallet.
    30  type Wallet struct {
    31  	// Version of the wallet, used for later upgrades.
    32  	Version string `json:"version"`
    33  
    34  	// A list of accounts which describes the details of each account
    35  	// in the wallet.
    36  	Accounts []*Account `json:"accounts"`
    37  
    38  	Scrypt keys.ScryptParams `json:"scrypt"`
    39  
    40  	// Extra metadata can be used for storing arbitrary data.
    41  	// This field can be empty.
    42  	Extra Extra `json:"extra"`
    43  
    44  	// Path where the wallet file is located..
    45  	path string
    46  }
    47  
    48  // Extra stores imported token contracts.
    49  type Extra struct {
    50  	// Tokens is a list of imported token contracts.
    51  	Tokens []*Token
    52  }
    53  
    54  // NewWallet creates a new NEO wallet at the given location.
    55  func NewWallet(location string) (*Wallet, error) {
    56  	file, err := os.Create(location)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	defer file.Close()
    61  	return newWallet(file), nil
    62  }
    63  
    64  // NewInMemoryWallet creates a new NEO wallet without linking to the read file on file system.
    65  // If wallet required to be written to the file system, [Wallet.SetPath] should be used to set the path.
    66  func NewInMemoryWallet() *Wallet {
    67  	return newWallet(nil)
    68  }
    69  
    70  // NewWalletFromFile creates a Wallet from the given wallet file path.
    71  func NewWalletFromFile(path string) (*Wallet, error) {
    72  	file, err := os.Open(path)
    73  	if err != nil {
    74  		return nil, fmt.Errorf("open wallet: %w", err)
    75  	}
    76  	defer file.Close()
    77  
    78  	wall := &Wallet{
    79  		path: file.Name(),
    80  	}
    81  	if err := json.NewDecoder(file).Decode(wall); err != nil {
    82  		return nil, fmt.Errorf("unmarshal wallet: %w", err)
    83  	}
    84  	return wall, nil
    85  }
    86  
    87  // NewWalletFromBytes creates a [Wallet] from the given byte slice.
    88  // Parameter wallet contains JSON representation of wallet, see [Wallet.JSON] for details.
    89  //
    90  // NewWalletFromBytes constructor doesn't set wallet's path. If you want to save the wallet to file system,
    91  // use [Wallet.SetPath].
    92  func NewWalletFromBytes(wallet []byte) (*Wallet, error) {
    93  	wall := &Wallet{}
    94  	if err := json.NewDecoder(bytes.NewReader(wallet)).Decode(wall); err != nil {
    95  		return nil, fmt.Errorf("unmarshal wallet: %w", err)
    96  	}
    97  
    98  	return wall, nil
    99  }
   100  
   101  func newWallet(rw io.ReadWriter) *Wallet {
   102  	var path string
   103  	if f, ok := rw.(*os.File); ok {
   104  		path = f.Name()
   105  	}
   106  	return &Wallet{
   107  		Version:  walletVersion,
   108  		Accounts: []*Account{},
   109  		Scrypt:   keys.NEP2ScryptParams(),
   110  		path:     path,
   111  	}
   112  }
   113  
   114  // CreateAccount generates a new account for the end user and encrypts
   115  // the private key with the given passphrase.
   116  func (w *Wallet) CreateAccount(name, passphrase string) error {
   117  	acc, err := NewAccount()
   118  	if err != nil {
   119  		return err
   120  	}
   121  	acc.Label = name
   122  	if err := acc.Encrypt(passphrase, w.Scrypt); err != nil {
   123  		return err
   124  	}
   125  	w.AddAccount(acc)
   126  	return w.Save()
   127  }
   128  
   129  // AddAccount adds an existing Account to the wallet.
   130  func (w *Wallet) AddAccount(acc *Account) {
   131  	w.Accounts = append(w.Accounts, acc)
   132  }
   133  
   134  // RemoveAccount removes an Account with the specified addr
   135  // from the wallet.
   136  func (w *Wallet) RemoveAccount(addr string) error {
   137  	for i, acc := range w.Accounts {
   138  		if acc.Address == addr {
   139  			copy(w.Accounts[i:], w.Accounts[i+1:])
   140  			w.Accounts = w.Accounts[:len(w.Accounts)-1]
   141  			return nil
   142  		}
   143  	}
   144  	return errors.New("account wasn't found")
   145  }
   146  
   147  // AddToken adds a new token to a wallet.
   148  func (w *Wallet) AddToken(tok *Token) {
   149  	w.Extra.Tokens = append(w.Extra.Tokens, tok)
   150  }
   151  
   152  // RemoveToken removes the token with the specified hash from the wallet.
   153  func (w *Wallet) RemoveToken(h util.Uint160) error {
   154  	for i, tok := range w.Extra.Tokens {
   155  		if tok.Hash.Equals(h) {
   156  			copy(w.Extra.Tokens[i:], w.Extra.Tokens[i+1:])
   157  			w.Extra.Tokens = w.Extra.Tokens[:len(w.Extra.Tokens)-1]
   158  			return nil
   159  		}
   160  	}
   161  	return errors.New("token wasn't found")
   162  }
   163  
   164  // Path returns the location of the wallet on the filesystem.
   165  func (w *Wallet) Path() string {
   166  	return w.path
   167  }
   168  
   169  // SetPath sets the location of the wallet on the filesystem.
   170  func (w *Wallet) SetPath(path string) {
   171  	w.path = path
   172  }
   173  
   174  // Save saves the wallet data to the file located at the path that was either provided
   175  // via [NewWalletFromFile] constructor or via [Wallet.SetPath].
   176  //
   177  // Returns [ErrPathIsEmpty] if wallet path is not set. See [Wallet.SetPath].
   178  func (w *Wallet) Save() error {
   179  	data, err := json.Marshal(w)
   180  	if err != nil {
   181  		return err
   182  	}
   183  
   184  	return w.writeRaw(data)
   185  }
   186  
   187  // SavePretty saves the wallet in a beautiful JSON.
   188  //
   189  // Returns [ErrPathIsEmpty] if wallet path is not set. See [Wallet.SetPath].
   190  func (w *Wallet) SavePretty() error {
   191  	data, err := json.MarshalIndent(w, "", "  ")
   192  	if err != nil {
   193  		return err
   194  	}
   195  
   196  	return w.writeRaw(data)
   197  }
   198  
   199  func (w *Wallet) writeRaw(data []byte) error {
   200  	if w.path == "" {
   201  		return ErrPathIsEmpty
   202  	}
   203  
   204  	return os.WriteFile(w.path, data, 0644)
   205  }
   206  
   207  // JSON outputs a pretty JSON representation of the wallet.
   208  func (w *Wallet) JSON() ([]byte, error) {
   209  	return json.MarshalIndent(w, " ", "	")
   210  }
   211  
   212  // Close closes all Wallet accounts making them incapable of signing anything
   213  // (unless they're decrypted again). It's not doing anything to the underlying
   214  // wallet file.
   215  func (w *Wallet) Close() {
   216  	for _, acc := range w.Accounts {
   217  		acc.Close()
   218  	}
   219  }
   220  
   221  // GetAccount returns an account corresponding to the provided scripthash.
   222  func (w *Wallet) GetAccount(h util.Uint160) *Account {
   223  	addr := address.Uint160ToString(h)
   224  	for _, acc := range w.Accounts {
   225  		if acc.Address == addr {
   226  			return acc
   227  		}
   228  	}
   229  
   230  	return nil
   231  }
   232  
   233  // GetChangeAddress returns the default address to send transaction's change to.
   234  func (w *Wallet) GetChangeAddress() util.Uint160 {
   235  	var res util.Uint160
   236  	var acc *Account
   237  
   238  	for i := range w.Accounts {
   239  		if acc == nil || w.Accounts[i].Default {
   240  			if w.Accounts[i].Contract != nil && vm.IsSignatureContract(w.Accounts[i].Contract.Script) {
   241  				acc = w.Accounts[i]
   242  				if w.Accounts[i].Default {
   243  					break
   244  				}
   245  			}
   246  		}
   247  	}
   248  	if acc != nil {
   249  		res = acc.Contract.ScriptHash()
   250  	}
   251  	return res
   252  }