<commit_before>/*
 * Copyright (c) 2013 Conformal Systems LLC <<EMAIL>>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package wallet

import (
	"bytes"
	"code.google.com/p/go.crypto/ripemd160"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/conformal/btcec"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
	"hash"
	"io"
	"math/big"
	"sync"
	"time"
)

const (
	// Length in bytes of KDF output.
	kdfOutputBytes = 32

	// Maximum length in bytes of a comment that can have a size represented
	// as a uint16.
	maxCommentLen = (1 << 16) - 1

	// Number of addresses to extend keypool by.
	nKeypoolIncrement = 100
)

const (
	defaultKdfComputeTime = 0.25
	defaultKdfMaxMem      = 32 * 1024 * 1024
)

// Possible errors when dealing with wallets.
var (
	ErrAddressNotFound    = errors.New("address not found")
	ErrChecksumMismatch   = errors.New("checksum mismatch")
	ErrMalformedEntry     = errors.New("malformed entry")
	ErrNetworkMismatch    = errors.New("network mismatch")
	ErrWalletDoesNotExist = errors.New("non-existant wallet")
	ErrWalletLocked       = errors.New("wallet is locked")
)

var (
	// '\xbaWALLET\x00'
	fileID = [8]byte{0xba, 0x57, 0x41, 0x4c, 0x4c, 0x45, 0x54, 0x00}

	mainnetMagicBytes = [4]byte{0xf9, 0xbe, 0xb4, 0xd9}
	testnetMagicBytes = [4]byte{0x0b, 0x11, 0x09, 0x07}
)

type entryHeader byte

const (
	addrCommentHeader entryHeader = 1 << iota
	txCommentHeader
	deletedHeader
	addrHeader entryHeader = 0
)

// We want to use binaryRead and binaryWrite instead of binary.Read
// and binary.Write because those from the binary package do not return
// the number of bytes actually written or read.  We need to return
// this value to correctly support the io.ReaderFrom and io.WriterTo
// interfaces.
func binaryRead(r io.Reader, order binary.ByteOrder, data interface{}) (n int64, err error) {
	var read int
	buf := make([]byte, binary.Size(data))
	if read, err = r.Read(buf); err != nil {
		return int64(read), err
	}
	if read < binary.Size(data) {
		return int64(read), io.EOF
	}
	return int64(read), binary.Read(bytes.NewBuffer(buf), order, data)
}

// See comment for binaryRead().
func binaryWrite(w io.Writer, order binary.ByteOrder, data interface{}) (n int64, err error) {
	var buf bytes.Buffer
	if err = binary.Write(&buf, order, data); err != nil {
		return 0, err
	}

	written, err := w.Write(buf.Bytes())
	return int64(written), err
}

// Calculate the hash of hasher over buf.
func calcHash(buf []byte, hasher hash.Hash) []byte {
	hasher.Write(buf)
	return hasher.Sum(nil)
}

// calculate hash160 which is ripemd160(sha256(data))
func calcHash160(buf []byte) []byte {
	return calcHash(calcHash(buf, sha256.New()), ripemd160.New())
}

// calculate hash256 which is sha256(sha256(data))
func calcHash256(buf []byte) []byte {
	return calcHash(calcHash(buf, sha256.New()), sha256.New())
}

// calculate sha512(data)
func calcSha512(buf []byte) []byte {
	return calcHash(buf, sha512.New())
}

// pubkeyFromPrivkey creates an encoded pubkey based on a
// 32-byte privkey.  The returned pubkey is 33 bytes if compressed,
// or 65 bytes if uncompressed.
func pubkeyFromPrivkey(privkey []byte, compress bool) (pubkey []byte) {
	x, y := btcec.S256().ScalarBaseMult(privkey)
	pub := (*btcec.PublicKey)(&ecdsa.PublicKey{
		Curve: btcec.S256(),
		X:     x,
		Y:     y,
	})

	if compress {
		return pub.SerializeCompressed()
	}
	return pub.SerializeUncompressed()
}

func keyOneIter(passphrase, salt []byte, memReqts uint64) []byte {
	saltedpass := append(passphrase, salt...)
	lutbl := make([]byte, memReqts)

	// Seed for lookup table
	seed := calcSha512(saltedpass)
	copy(lutbl[:sha512.Size], seed)

	for nByte := 0; nByte < (int(memReqts) - sha512.Size); nByte += sha512.Size {
		hash := calcSha512(lutbl[nByte : nByte+sha512.Size])
		copy(lutbl[nByte+sha512.Size:nByte+2*sha512.Size], hash[:])
	}

	x := lutbl[cap(lutbl)-sha512.Size:]

	seqCt := uint32(memReqts / sha512.Size)
	nLookups := seqCt / 2
	for i := uint32(0); i < nLookups; i++ {
		// Armory ignores endianness here.  We assume LE.
		newIdx := binary.LittleEndian.Uint32(x[cap(x)-4:]) % seqCt

		// Index of hash result at newIdx
		vIdx := newIdx * sha512.Size
		v := lutbl[vIdx : vIdx+sha512.Size]

		// XOR hash x with hash v
		for j := 0; j < sha512.Size; j++ {
			x[j] ^= v[j]
		}

		// Save new hash to x
		hash := calcSha512(x)
		copy(x, hash[:])
	}

	return x[:kdfOutputBytes]
}

// Key implements the key derivation function used by Armory
// based on the ROMix algorithm described in <NAME>'s paper
// "Stronger Key Derivation via Sequential Memory-Hard Functions"
// (http://www.tarsnap.com/scrypt/scrypt.pdf).
func Key(passphrase []byte, params *kdfParameters) []byte {
	masterKey := passphrase
	for i := uint32(0); i < params.nIter; i++ {
		masterKey = keyOneIter(masterKey, params.salt[:], params.mem)
	}
	return masterKey
}

func pad(size int, b []byte) []byte {
	// Prevent a possible panic if the input exceeds the expected size.
	if len(b) > size {
		size = len(b)
	}

	p := make([]byte, size)
	copy(p[size-len(b):], b)
	return p
}

// ChainedPrivKey deterministically generates a new private key using a
// previous address and chaincode.  privkey and chaincode must be 32
// bytes long, and pubkey may either be 33 bytes, 65 bytes or nil (in
// which case it is generated by the privkey).
func ChainedPrivKey(privkey, pubkey, chaincode []byte) ([]byte, error) {
	if len(privkey) != 32 {
		return nil, fmt.Errorf("invalid privkey length %d (must be 32)",
			len(privkey))
	}
	if len(chaincode) != 32 {
		return nil, fmt.Errorf("invalid chaincode length %d (must be 32)",
			len(chaincode))
	}
	if pubkey == nil {
		pubkey = pubkeyFromPrivkey(privkey, true)
	} else if !(len(pubkey) == 65 || len(pubkey) == 33) {
		return nil, fmt.Errorf("invalid pubkey length %d", len(pubkey))
	}

	// This is a perfect example of YOLO crypto.  Armory claims this XORing
	// with the SHA256 hash of the pubkey is done to add extra entropy (why
	// you'd want to add entropy to a deterministic function, I don't know),
	// even though the pubkey is generated directly from the privkey.  In
	// terms of security or privacy, this is a complete waste of CPU cycles,
	// but we do the same because we want to keep compatibility with
	// Armory's chained address generation.
	xorbytes := make([]byte, 32)
	chainMod := calcHash256(pubkey)
	for i := range xorbytes {
		xorbytes[i] = chainMod[i] ^ chaincode[i]
	}
	chainXor := new(big.Int).SetBytes(xorbytes)
	privint := new(big.Int).SetBytes(privkey)

	t := new(big.Int).Mul(chainXor, privint)
	b := t.Mod(t, btcec.S256().N).Bytes()
	return pad(32, b), nil
}

type varEntries []io.WriterTo

func (v *varEntries) WriteTo(w io.Writer) (n int64, err error) {
	ss := []io.WriterTo(*v)

	var written int64
	for _, s := range ss {
		var err error
		if written, err = s.WriteTo(w); err != nil {
			return n + written, err
		}
		n += written
	}
	return n, nil
}

func (v *varEntries) ReadFrom(r io.Reader) (n int64, err error) {
	var read int64

	// Remove any previous entries.
	*v = nil
	wts := []io.WriterTo(*v)

	// Keep reading entries until an EOF is reached.
	for {
		var header entryHeader
		if read, err = binaryRead(r, binary.LittleEndian, &header); err != nil {
			// EOF here is not an error.
			if err == io.EOF {
				return n + read, nil
			}
			return n + read, err
		}
		n += read

		var wt io.WriterTo
		switch header {
		case addrHeader:
			var entry addrEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
			wt = &entry
		case addrCommentHeader:
			var entry addrCommentEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
			wt = &entry
		case txCommentHeader:
			var entry txCommentEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
			wt = &entry
		case deletedHeader:
			var entry deletedEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
		default:
			return n, fmt.Errorf("unknown entry header: %d", uint8(header))
		}
		if wt != nil {
			wts = append(wts, wt)
			*v = wts
		}
	}
}

type addressHashKey string
type transactionHashKey string
type comment []byte

// Wallet represents an btcd/Armory wallet in memory.  It
// implements the io.ReaderFrom and io.WriterTo interfaces to read
// from and write to any type of byte streams, including files.
type Wallet struct {
	version      uint32
	net          btcwire.BitcoinNet
	flags        walletFlags
	uniqID       [6]byte
	createDate   int64
	name         [32]byte
	desc         [256]byte
	highestUsed  int64
	kdfParams    kdfParameters
	keyGenerator btcAddress

	// These are non-standard and fit in the extra 1024 bytes between the
	// root address and the appended entries.
	syncedBlockHeight int32
	syncedBlockHash   btcwire.ShaHash

	addrMap        map[addressHashKey]*btcAddress
	addrCommentMap map[addressHashKey]comment
	txCommentMap   map[transactionHashKey]comment

	// These are not serialized.
	secret struct {
		sync.Mutex
		key []byte
	}
	chainIdxMap   map[int64]addressHashKey
	importedAddrs []*btcAddress
	lastChainIdx  int64
}

// UnusedWalletBytes specifies the number of actually unused bytes
// between the root address and the appended entries in a serialized
// wallet.  Armory's wallet file format provides 1024 unused bytes
// in this space.  btcwallet requires saving a few additional details
// with the wallet file, so the binary sizes of those are subtracted
// from 1024.  Currently, these are:
//
//  - last synced block height (int32, 4 bytes)
//  - last synced block hash (btcwire.ShaHash, btcwire.HashSize bytes)
const UnusedWalletBytes = 1024 - 4 - btcwire.HashSize

// NewWallet creates and initializes a new Wallet.  name's and
// desc's binary representation must not exceed 32 and 256 bytes,
// respectively.  All address private keys are encrypted with passphrase.
// The wallet is returned unlocked.
func NewWallet(name, desc string, passphrase []byte, net btcwire.BitcoinNet,
	createdAt *BlockStamp) (*Wallet, error) {

	// Check sizes of inputs.
	if len([]byte(name)) > 32 {
		return nil, errors.New("name exceeds 32 byte maximum size")
	}
	if len([]byte(desc)) > 256 {
		return nil, errors.New("desc exceeds 256 byte maximum size")
	}

	// Randomly-generate rootkey and chaincode.
	rootkey, chaincode := make([]byte, 32), make([]byte, 32)
	if _, err := rand.Read(rootkey); err != nil {
		return nil, err
	}
	if _, err := rand.Read(chaincode); err != nil {
		return nil, err
	}

	// Create new root address from key and chaincode.
	root, err := newRootBtcAddress(rootkey, nil, chaincode, createdAt)
	if err != nil {
		return nil, err
	}

	// Compute AES key and encrypt root address.
	kdfp, err := computeKdfParameters(defaultKdfComputeTime, defaultKdfMaxMem)
	if err != nil {
		return nil, err
	}
	aeskey := Key([]byte(passphrase), kdfp)
	if err := root.encrypt(aeskey); err != nil {
		return nil, err
	}

	// Create and fill wallet.
	w := &Wallet{
		version: 0, // TODO(jrick): implement versioning
		// TODO(jrick): not sure we will need uniqID, but would be good for
		// compat with armory.
		net: net,
		flags: walletFlags{
			useEncryption: true,
			watchingOnly:  false,
		},
		createDate:        time.Now().Unix(),
		highestUsed:       rootKeyChainIdx,
		kdfParams:         *kdfp,
		keyGenerator:      *root,
		syncedBlockHeight: createdAt.Height,
		syncedBlockHash:   createdAt.Hash,
		addrMap:           make(map[addressHashKey]*btcAddress),
		addrCommentMap:    make(map[addressHashKey]comment),
		txCommentMap:      make(map[transactionHashKey]comment),
		chainIdxMap:       make(map[int64]addressHashKey),
		lastChainIdx:      rootKeyChainIdx,
	}
	copy(w.name[:], []byte(name))
	copy(w.desc[:], []byte(desc))

	// Add root address to maps.
	w.addrMap[addressHashKey(w.keyGenerator.pubKeyHash[:])] = &w.keyGenerator
	w.chainIdxMap[rootKeyChainIdx] = addressHashKey(w.keyGenerator.pubKeyHash[:])

	// Fill keypool.
	if err := w.extendKeypool(nKeypoolIncrement, aeskey, createdAt); err != nil {
		return nil, err
	}

	return w, nil
}

// Name returns the name of a wallet.  This name is used as the
// account name for btcwallet JSON methods.
func (w *Wallet) Name() string {
	return string(w.name[:])
}

// ReadFrom reads data from a io.Reader and saves it to a Wallet,
// returning the number of bytes read and any errors encountered.
func (w *Wallet) ReadFrom(r io.Reader) (n int64, err error) {
	var read int64

	w.addrMap = make(map[addressHashKey]*btcAddress)
	w.addrCommentMap = make(map[addressHashKey]comment)
	w.chainIdxMap = make(map[int64]addressHashKey)
	w.txCommentMap = make(map[transactionHashKey]comment)

	var id [8]byte
	var appendedEntries varEntries

	// Iterate through each entry needing to be read.  If data
	// implements io.ReaderFrom, use its ReadFrom func.  Otherwise,
	// data is a pointer to a fixed sized value.
	datas := []interface{}{
		&id,<commit_msg>Ignore zeros when generating wallet name.<commit_after>/*
 * Copyright (c) 2013 Conformal Systems LLC <<EMAIL>>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package wallet

import (
	"bytes"
	"code.google.com/p/go.crypto/ripemd160"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/conformal/btcec"
	"github.com/conformal/btcutil"
	"github.com/conformal/btcwire"
	"hash"
	"io"
	"math/big"
	"sync"
	"time"
)

const (
	// Length in bytes of KDF output.
	kdfOutputBytes = 32

	// Maximum length in bytes of a comment that can have a size represented
	// as a uint16.
	maxCommentLen = (1 << 16) - 1

	// Number of addresses to extend keypool by.
	nKeypoolIncrement = 100
)

const (
	defaultKdfComputeTime = 0.25
	defaultKdfMaxMem      = 32 * 1024 * 1024
)

// Possible errors when dealing with wallets.
var (
	ErrAddressNotFound    = errors.New("address not found")
	ErrChecksumMismatch   = errors.New("checksum mismatch")
	ErrMalformedEntry     = errors.New("malformed entry")
	ErrNetworkMismatch    = errors.New("network mismatch")
	ErrWalletDoesNotExist = errors.New("non-existant wallet")
	ErrWalletLocked       = errors.New("wallet is locked")
)

var (
	// '\xbaWALLET\x00'
	fileID = [8]byte{0xba, 0x57, 0x41, 0x4c, 0x4c, 0x45, 0x54, 0x00}

	mainnetMagicBytes = [4]byte{0xf9, 0xbe, 0xb4, 0xd9}
	testnetMagicBytes = [4]byte{0x0b, 0x11, 0x09, 0x07}
)

type entryHeader byte

const (
	addrCommentHeader entryHeader = 1 << iota
	txCommentHeader
	deletedHeader
	addrHeader entryHeader = 0
)

// We want to use binaryRead and binaryWrite instead of binary.Read
// and binary.Write because those from the binary package do not return
// the number of bytes actually written or read.  We need to return
// this value to correctly support the io.ReaderFrom and io.WriterTo
// interfaces.
func binaryRead(r io.Reader, order binary.ByteOrder, data interface{}) (n int64, err error) {
	var read int
	buf := make([]byte, binary.Size(data))
	if read, err = r.Read(buf); err != nil {
		return int64(read), err
	}
	if read < binary.Size(data) {
		return int64(read), io.EOF
	}
	return int64(read), binary.Read(bytes.NewBuffer(buf), order, data)
}

// See comment for binaryRead().
func binaryWrite(w io.Writer, order binary.ByteOrder, data interface{}) (n int64, err error) {
	var buf bytes.Buffer
	if err = binary.Write(&buf, order, data); err != nil {
		return 0, err
	}

	written, err := w.Write(buf.Bytes())
	return int64(written), err
}

// Calculate the hash of hasher over buf.
func calcHash(buf []byte, hasher hash.Hash) []byte {
	hasher.Write(buf)
	return hasher.Sum(nil)
}

// calculate hash160 which is ripemd160(sha256(data))
func calcHash160(buf []byte) []byte {
	return calcHash(calcHash(buf, sha256.New()), ripemd160.New())
}

// calculate hash256 which is sha256(sha256(data))
func calcHash256(buf []byte) []byte {
	return calcHash(calcHash(buf, sha256.New()), sha256.New())
}

// calculate sha512(data)
func calcSha512(buf []byte) []byte {
	return calcHash(buf, sha512.New())
}

// pubkeyFromPrivkey creates an encoded pubkey based on a
// 32-byte privkey.  The returned pubkey is 33 bytes if compressed,
// or 65 bytes if uncompressed.
func pubkeyFromPrivkey(privkey []byte, compress bool) (pubkey []byte) {
	x, y := btcec.S256().ScalarBaseMult(privkey)
	pub := (*btcec.PublicKey)(&ecdsa.PublicKey{
		Curve: btcec.S256(),
		X:     x,
		Y:     y,
	})

	if compress {
		return pub.SerializeCompressed()
	}
	return pub.SerializeUncompressed()
}

func keyOneIter(passphrase, salt []byte, memReqts uint64) []byte {
	saltedpass := append(passphrase, salt...)
	lutbl := make([]byte, memReqts)

	// Seed for lookup table
	seed := calcSha512(saltedpass)
	copy(lutbl[:sha512.Size], seed)

	for nByte := 0; nByte < (int(memReqts) - sha512.Size); nByte += sha512.Size {
		hash := calcSha512(lutbl[nByte : nByte+sha512.Size])
		copy(lutbl[nByte+sha512.Size:nByte+2*sha512.Size], hash[:])
	}

	x := lutbl[cap(lutbl)-sha512.Size:]

	seqCt := uint32(memReqts / sha512.Size)
	nLookups := seqCt / 2
	for i := uint32(0); i < nLookups; i++ {
		// Armory ignores endianness here.  We assume LE.
		newIdx := binary.LittleEndian.Uint32(x[cap(x)-4:]) % seqCt

		// Index of hash result at newIdx
		vIdx := newIdx * sha512.Size
		v := lutbl[vIdx : vIdx+sha512.Size]

		// XOR hash x with hash v
		for j := 0; j < sha512.Size; j++ {
			x[j] ^= v[j]
		}

		// Save new hash to x
		hash := calcSha512(x)
		copy(x, hash[:])
	}

	return x[:kdfOutputBytes]
}

// Key implements the key derivation function used by Armory
// based on the ROMix algorithm described in <NAME>'s paper
// "Stronger Key Derivation via Sequential Memory-Hard Functions"
// (http://www.tarsnap.com/scrypt/scrypt.pdf).
func Key(passphrase []byte, params *kdfParameters) []byte {
	masterKey := passphrase
	for i := uint32(0); i < params.nIter; i++ {
		masterKey = keyOneIter(masterKey, params.salt[:], params.mem)
	}
	return masterKey
}

func pad(size int, b []byte) []byte {
	// Prevent a possible panic if the input exceeds the expected size.
	if len(b) > size {
		size = len(b)
	}

	p := make([]byte, size)
	copy(p[size-len(b):], b)
	return p
}

// ChainedPrivKey deterministically generates a new private key using a
// previous address and chaincode.  privkey and chaincode must be 32
// bytes long, and pubkey may either be 33 bytes, 65 bytes or nil (in
// which case it is generated by the privkey).
func ChainedPrivKey(privkey, pubkey, chaincode []byte) ([]byte, error) {
	if len(privkey) != 32 {
		return nil, fmt.Errorf("invalid privkey length %d (must be 32)",
			len(privkey))
	}
	if len(chaincode) != 32 {
		return nil, fmt.Errorf("invalid chaincode length %d (must be 32)",
			len(chaincode))
	}
	if pubkey == nil {
		pubkey = pubkeyFromPrivkey(privkey, true)
	} else if !(len(pubkey) == 65 || len(pubkey) == 33) {
		return nil, fmt.Errorf("invalid pubkey length %d", len(pubkey))
	}

	// This is a perfect example of YOLO crypto.  Armory claims this XORing
	// with the SHA256 hash of the pubkey is done to add extra entropy (why
	// you'd want to add entropy to a deterministic function, I don't know),
	// even though the pubkey is generated directly from the privkey.  In
	// terms of security or privacy, this is a complete waste of CPU cycles,
	// but we do the same because we want to keep compatibility with
	// Armory's chained address generation.
	xorbytes := make([]byte, 32)
	chainMod := calcHash256(pubkey)
	for i := range xorbytes {
		xorbytes[i] = chainMod[i] ^ chaincode[i]
	}
	chainXor := new(big.Int).SetBytes(xorbytes)
	privint := new(big.Int).SetBytes(privkey)

	t := new(big.Int).Mul(chainXor, privint)
	b := t.Mod(t, btcec.S256().N).Bytes()
	return pad(32, b), nil
}

type varEntries []io.WriterTo

func (v *varEntries) WriteTo(w io.Writer) (n int64, err error) {
	ss := []io.WriterTo(*v)

	var written int64
	for _, s := range ss {
		var err error
		if written, err = s.WriteTo(w); err != nil {
			return n + written, err
		}
		n += written
	}
	return n, nil
}

func (v *varEntries) ReadFrom(r io.Reader) (n int64, err error) {
	var read int64

	// Remove any previous entries.
	*v = nil
	wts := []io.WriterTo(*v)

	// Keep reading entries until an EOF is reached.
	for {
		var header entryHeader
		if read, err = binaryRead(r, binary.LittleEndian, &header); err != nil {
			// EOF here is not an error.
			if err == io.EOF {
				return n + read, nil
			}
			return n + read, err
		}
		n += read

		var wt io.WriterTo
		switch header {
		case addrHeader:
			var entry addrEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
			wt = &entry
		case addrCommentHeader:
			var entry addrCommentEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
			wt = &entry
		case txCommentHeader:
			var entry txCommentEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
			wt = &entry
		case deletedHeader:
			var entry deletedEntry
			if read, err = entry.ReadFrom(r); err != nil {
				return n + read, err
			}
			n += read
		default:
			return n, fmt.Errorf("unknown entry header: %d", uint8(header))
		}
		if wt != nil {
			wts = append(wts, wt)
			*v = wts
		}
	}
}

type addressHashKey string
type transactionHashKey string
type comment []byte

// Wallet represents an btcd/Armory wallet in memory.  It
// implements the io.ReaderFrom and io.WriterTo interfaces to read
// from and write to any type of byte streams, including files.
type Wallet struct {
	version      uint32
	net          btcwire.BitcoinNet
	flags        walletFlags
	uniqID       [6]byte
	createDate   int64
	name         [32]byte
	desc         [256]byte
	highestUsed  int64
	kdfParams    kdfParameters
	keyGenerator btcAddress

	// These are non-standard and fit in the extra 1024 bytes between the
	// root address and the appended entries.
	syncedBlockHeight int32
	syncedBlockHash   btcwire.ShaHash

	addrMap        map[addressHashKey]*btcAddress
	addrCommentMap map[addressHashKey]comment
	txCommentMap   map[transactionHashKey]comment

	// These are not serialized.
	secret struct {
		sync.Mutex
		key []byte
	}
	chainIdxMap   map[int64]addressHashKey
	importedAddrs []*btcAddress
	lastChainIdx  int64
}

// UnusedWalletBytes specifies the number of actually unused bytes
// between the root address and the appended entries in a serialized
// wallet.  Armory's wallet file format provides 1024 unused bytes
// in this space.  btcwallet requires saving a few additional details
// with the wallet file, so the binary sizes of those are subtracted
// from 1024.  Currently, these are:
//
//  - last synced block height (int32, 4 bytes)
//  - last synced block hash (btcwire.ShaHash, btcwire.HashSize bytes)
const UnusedWalletBytes = 1024 - 4 - btcwire.HashSize

// NewWallet creates and initializes a new Wallet.  name's and
// desc's binary representation must not exceed 32 and 256 bytes,
// respectively.  All address private keys are encrypted with passphrase.
// The wallet is returned unlocked.
func NewWallet(name, desc string, passphrase []byte, net btcwire.BitcoinNet,
	createdAt *BlockStamp) (*Wallet, error) {

	// Check sizes of inputs.
	if len([]byte(name)) > 32 {
		return nil, errors.New("name exceeds 32 byte maximum size")
	}
	if len([]byte(desc)) > 256 {
		return nil, errors.New("desc exceeds 256 byte maximum size")
	}

	// Randomly-generate rootkey and chaincode.
	rootkey, chaincode := make([]byte, 32), make([]byte, 32)
	if _, err := rand.Read(rootkey); err != nil {
		return nil, err
	}
	if _, err := rand.Read(chaincode); err != nil {
		return nil, err
	}

	// Create new root address from key and chaincode.
	root, err := newRootBtcAddress(rootkey, nil, chaincode, createdAt)
	if err != nil {
		return nil, err
	}

	// Compute AES key and encrypt root address.
	kdfp, err := computeKdfParameters(defaultKdfComputeTime, defaultKdfMaxMem)
	if err != nil {
		return nil, err
	}
	aeskey := Key([]byte(passphrase), kdfp)
	if err := root.encrypt(aeskey); err != nil {
		return nil, err
	}

	// Create and fill wallet.
	w := &Wallet{
		version: 0, // TODO(jrick): implement versioning
		// TODO(jrick): not sure we will need uniqID, but would be good for
		// compat with armory.
		net: net,
		flags: walletFlags{
			useEncryption: true,
			watchingOnly:  false,
		},
		createDate:        time.Now().Unix(),
		highestUsed:       rootKeyChainIdx,
		kdfParams:         *kdfp,
		keyGenerator:      *root,
		syncedBlockHeight: createdAt.Height,
		syncedBlockHash:   createdAt.Hash,
		addrMap:           make(map[addressHashKey]*btcAddress),
		addrCommentMap:    make(map[addressHashKey]comment),
		txCommentMap:      make(map[transactionHashKey]comment),
		chainIdxMap:       make(map[int64]addressHashKey),
		lastChainIdx:      rootKeyChainIdx,
	}
	copy(w.name[:], []byte(name))
	copy(w.desc[:], []byte(desc))

	// Add root address to maps.
	w.addrMap[addressHashKey(w.keyGenerator.pubKeyHash[:])] = &w.keyGenerator
	w.chainIdxMap[rootKeyChainIdx] = addressHashKey(w.keyGenerator.pubKeyHash[:])

	// Fill keypool.
	if err := w.extendKeypool(nKeypoolIncrement, aeskey, createdAt); err != nil {
		return nil, err
	}

	return w, nil
}

// Name returns the name of a wallet.  This name is used as the
// account name for btcwallet JSON methods.
func (w *Wallet) Name() string {
	last := len(w.name[:])
	for i, b := range w.name[:] {
		if b == 0x00 {
			last = i
			break
		}
	}
	return string(w.name[:last])
}

// ReadFrom reads data from a io.Reader and saves it to a Wallet,
// returning the number of bytes read and any errors encountered.
func (w *Wallet) ReadFrom(r io.Reader) (n int64, err error) {
	var read int64

	w.addrMap = make(map[addressHashKey]*btcAddress)
	w.addrCommentMap = make(map[addressHashKey]comment)
	w.chainIdxMap = make(map[int64]addressHashKey)
	w.txCommentMap = make(map[transactionHashKey]comment)

	var id [8]byte
	var appendedEntries varEntries

	// Iterate through each entry needing to be read.  If data
	// implements io.ReaderFrom, use its ReadFrom func.  Otherwise,
	// data is a pointer to a fixed sized value.
	datas := []interface{}{
		&id,
{'fraction_non_alphanumeric': 0.11687418560880267, 'fraction_numerical': 0.023128710004343418, 'mean_word_length': 3.624100418410042, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '759215', 'n_tokens_mistral': 10244, 'n_tokens_neox': 9270, 'n_words': 3073}