Allow to create multiple wallets

cli.go

@@ -18,14 +18,11 @@ func (cli *CLI) createBlockchain(address string) {
 }
 
 func (cli *CLI) createWallet() {
-	wallet, err := NewWallet()
-	if err == nil {
-		fmt.Printf("Your address: %s\n", wallet.GetAddress())
-		wallet.SaveToFile()
-	} else {
-		fmt.Println(err.Error())
-		os.Exit(1)
-	}
+	wallets := NewWallets()
+	address := wallets.CreateWallet()
+	wallets.SaveToFile()
+
+	fmt.Printf("Your new address: %s\n", address)
 }
 
 func (cli *CLI) getBalance(address string) {

wallet.go

@@ -7,7 +7,7 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"encoding/gob"
-	"errors"
+	"fmt"
 	"io/ioutil"
 	"log"
 	"os"
@@ -18,12 +18,17 @@ import (
 const version = byte(0x00)
 const walletFile = "wallet.dat"
 
-// Wallet ...
+// Wallet stores private and public keys
 type Wallet struct {
 	PrivateKey ecdsa.PrivateKey
 	PublicKey  ecdsa.PublicKey
 }
 
+// Wallets stores a collection of wallets
+type Wallets struct {
+	Wallets map[string]*Wallet
+}
+
 // GetAddress returns wallet address
 func (w Wallet) GetAddress() []byte {
 	public := append(w.PublicKey.X.Bytes(), w.PublicKey.Y.Bytes()...)
@@ -45,32 +50,12 @@ func (w Wallet) GetAddress() []byte {
 	return address
 }
 
-// SaveToFile saves the wallet to a file
-func (w Wallet) SaveToFile() {
-	var content bytes.Buffer
-
-	gob.Register(w.PrivateKey.Curve)
-	encoder := gob.NewEncoder(&content)
-	err := encoder.Encode(w)
-	if err != nil {
-		log.Panic(err)
-	}
-
-	err = ioutil.WriteFile(walletFile, content.Bytes(), 0644)
-	if err != nil {
-		log.Panic(err)
-	}
-}
-
 // NewWallet creates and returns a Wallet
-func NewWallet() (*Wallet, error) {
-	if _, err := os.Stat(walletFile); !os.IsNotExist(err) {
-		return nil, errors.New("Wallet already exists")
-	}
+func NewWallet() *Wallet {
 	private, public := newKeyPair()
 	wallet := Wallet{private, public}
 
-	return &wallet, nil
+	return &wallet
 }
 
 func newKeyPair() (ecdsa.PrivateKey, ecdsa.PublicKey) {
@@ -83,6 +68,73 @@ func newKeyPair() (ecdsa.PrivateKey, ecdsa.PublicKey) {
 	return *private, private.PublicKey
 }
 
+// CreateWallet adds a Wallet to Wallets
+func (ws *Wallets) CreateWallet() string {
+	wallet := NewWallet()
+	address := fmt.Sprintf("%s", wallet.GetAddress())
+
+	ws.Wallets[address] = wallet
+
+	return address
+}
+
+// SaveToFile saves wallets to a file
+func (ws Wallets) SaveToFile() {
+	var content bytes.Buffer
+
+	gob.Register(elliptic.P256())
+
+	encoder := gob.NewEncoder(&content)
+	err := encoder.Encode(ws)
+	if err != nil {
+		log.Panic(err)
+	}
+
+	err = ioutil.WriteFile(walletFile, content.Bytes(), 0644)
+	if err != nil {
+		log.Panic(err)
+	}
+}
+
+// LoadFromFile loads wallets from the file
+func (ws *Wallets) LoadFromFile() error {
+	if _, err := os.Stat(walletFile); os.IsNotExist(err) {
+		return err
+	}
+
+	fileContent, err := ioutil.ReadFile(walletFile)
+	if err != nil {
+		log.Panic(err)
+	}
+
+	var wallets Wallets
+	gob.Register(elliptic.P256())
+	decoder := gob.NewDecoder(bytes.NewReader(fileContent))
+	err = decoder.Decode(&wallets)
+	if err != nil {
+		log.Panic(err)
+	}
+
+	ws.Wallets = wallets.Wallets
+
+	return nil
+}
+
+// NewWallets ...
+func NewWallets() *Wallets {
+	wallets := Wallets{}
+	wallets.Wallets = make(map[string]*Wallet)
+
+	err := wallets.LoadFromFile()
+	if err != nil {
+		fmt.Println("Wallets file doesn't exist")
+		// wallets.CreateWallet()
+		// wallets.SaveToFile()
+	}
+
+	return &wallets
+}
+
 // Checksum ...
 func checksum(payload []byte) []byte {
 	firstSHA := sha256.Sum256(payload)