# addressset.py -- btcrecover AddressSet library
# Copyright (C) 2017 Christopher Gurnee
#
# This file is part of btcrecover.
#
# btcrecover is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version
# 2 of the License, or (at your option) any later version.
#
# btcrecover is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see http://www.gnu.org/licenses/

# If you find this program helpful, please consider a small
# donation to the developer at the following Bitcoin address:
#
#           3Au8ZodNHPei7MQiSVAWb7NB2yqsb48GW4
#
#                      Thank You!

# (all optional futures for 2.7 except unicode_literals)
from __future__ import print_function, absolute_import, division

__version__ =  "0.1.3"

import struct, base64, io, mmap, ast, itertools, sys, gc, glob
from pybitcointools import bin_hash160, decompress
from os import path


def bytes_to_int(bytes_rep):
    """convert a string of bytes (in big-endian order) to an integer

    :param bytes_rep: the raw bytes
    :type bytes_rep: str
    :return: the unsigned integer
    :rtype: int or long
    """
    bytes_len = len(bytes_rep)
    if bytes_len <= 4:
        return struct.unpack(">I", (4-bytes_len)*b"\0" + bytes_rep)[0]
    return long(base64.b16encode(bytes_rep), 16)


class AddressSet(object):
    """
    A set-like collection optimized for testing membership of Bitcoin addresses
    from their raw hash160 format with support for serializing to/from files
    """
    VERSION    = 1
    MAGIC      = b"seedrecover address database\r\n"  # file magic
    HEADER_LEN = 65536
    assert HEADER_LEN % mmap.ALLOCATIONGRANULARITY == 0

    def __init__(self, table_len, bytes_per_addr = 8, max_load = 0.75):
        """
        :param table_len: hash table size in count of addresses; must be a power of 2
        :type table_len: int
        :param bytes_per_addr: number of bytes of each address to store in the hash table
        :type bytes_per_addr: int
        :param max_load: max permissible load factor before an exception is raised
        :type max_load: float
        """
        if table_len < 1 or 1 << (table_len.bit_length()-1) != table_len:
            raise ValueError("table_len must be a positive power of 2")
        if not 1 <= bytes_per_addr <= 19:
            raise ValueError("bytes_per_addr must be between 1 and 19 inclusive")
        if not 0.0 < max_load < 1.0:
            raise ValueError("max_load must be between 0.0 and 1.0 exclusive")
        self._table_bytes    = table_len * bytes_per_addr         # len of hash table in bytes
        self._bytes_per_addr = bytes_per_addr                     # number of bytes per address to store
        self._null_addr      = b"\0" * bytes_per_addr             # all 0s is an empty hash table slot
        self._len            = 0                                  # count of addresses in the set
        self._max_len        = int(table_len * max_load)          # beyond this violates the load factor
        self._hash_bytes     = (table_len.bit_length() + 6) // 8  # number of bytes required for the mask
        self._hash_mask      = table_len - 1                      # mask used for the hash function
        self._data           = bytearray(self._table_bytes)       # the table itself
        self._dbfile         = None                               # file object, its .name is req'd for pickling
        self._mmap_access    = None                               # also required for pickling
        self.last_filenum    = None                               # will be serialized if set by the user
        if self._bytes_per_addr + self._hash_bytes > 20:
            raise ValueError("not enough bytes for both hashing and storage; "
                             "reduce either the bytes_per_addr or table_len")

    def __getstate__(self):
        # mmaps can't be pickled, so save only what's needed to recreate the object from scratch later
        if isinstance(self._data, mmap.mmap):
            return {"dbfilename": self._dbfile.name, "mmap_access": self._mmap_access}
        else:
            return self.__dict__

    def __setstate__(self, state):
        # If the object contained an mmap, recreate it from scratch
        if "dbfilename" in state:
            new = self.fromfile(open(state["dbfilename"], "r+b" if state["mmap_access"]==mmap.ACCESS_WRITE else "rb"),
                                mmap_access=state["mmap_access"], preload=False)
            self.__dict__ = new.__dict__.copy()
            new._dbfile = new._data = None  # ensure new's __del__() doesn't close() anything
        else:
            self.__dict__ = state

    def  __len__(self):
        return self._len

    def __contains__(self, address):
        return self._find(address) is True

    def add(self, address):
        """Adds the address to the set

        :param address: the address in hash160 (length 20) format to add
        :type address: bytes or str
        """
        pos = self._find(address)
        if pos is not True:
            bytes_to_add = address[ -(self._bytes_per_addr+self._hash_bytes) : -self._hash_bytes]
            if bytes_to_add.endswith(self._null_addr):
                return  # ignore these invalid addresses
            if self._len >= self._max_len:
                raise ValueError("addition to AddressSet exceeds load factor")
            self._data[pos : pos+self._bytes_per_addr] = bytes_to_add
            self._len += 1

    # Hash table with open addressing and linear probing:
    # The hash function is simply some of the address's least significant bits (since
    # most addresses are random hashes, this and linear probing should be sufficient).
    # To further save space, only the least significant _bytes_per_addr (typ. 8) bytes--
    # excluding those bytes already used for the "hash" above--are stored in the table,
    # causing different addresses to appear to be the same and false positives, however
    # (with high probability) only for invalid addresses (those w/o private keys).
    def _find(self, addr_to_find):
        pos = self._bytes_per_addr * (bytes_to_int(addr_to_find[ -self._hash_bytes :]) & self._hash_mask)
        while True:
            cur_addr = self._data[pos : pos+self._bytes_per_addr]
            if cur_addr == self._null_addr:
                return pos  # the position this element could be inserted at
            if len(addr_to_find) > self._bytes_per_addr:
                addr_to_find = addr_to_find[ -(self._bytes_per_addr+self._hash_bytes) : -self._hash_bytes]
            if cur_addr == addr_to_find:
                return True
            pos += self._bytes_per_addr  # linear probing
            if pos >= self._table_bytes:
                pos = 0

    def __iter__(self):
        """Iterates over the set returning the bytes_per_addr stored for each address
        """
        pos = 0
        while pos < self._table_bytes:
            cur_addr = self._data[pos : pos+self._bytes_per_addr]
            if cur_addr != self._null_addr:
                yield cur_addr
            pos += self._bytes_per_addr

    def __reversed__(self):
        pos = self._table_bytes - self._bytes_per_addr
        while pos >= 0:
            cur_addr = self._data[pos : pos+self._bytes_per_addr]
            if cur_addr != self._null_addr:
                yield cur_addr
            pos -= self._bytes_per_addr

    @staticmethod
    def _remove_nonheader_attribs(attrs):
        del attrs["_data"], attrs["_dbfile"], attrs["_mmap_access"]

    def _header(self):
        # Construct a 64K header with the file magic, this object's attributes, plus the version
        header_dict = self.__dict__.copy()
        self._remove_nonheader_attribs(header_dict)
        header_dict["version"] = self.VERSION
        header = repr(header_dict) + b"\r\n"
        assert ast.literal_eval(header) == header_dict
        header = self.MAGIC + header
        header_len = len(header)
        assert header_len < self.HEADER_LEN
        return header + b"\0" * (self.HEADER_LEN - header_len)  # appends at least one nul

    def tofile(self, dbfile):
        """Save the address set to a file

        :param dbfile: an open file object where the set is saved (overwriting it)
        :type dbfile: io.FileIO or file
        """
        if dbfile.tell() % mmap.ALLOCATIONGRANULARITY != 0:
            print("AddressSet: warning: if header position in file isn't a multiple of {}, it probably can't be loaded with fromfile()"
                  .format(mmap.ALLOCATIONGRANULARITY), file=sys.stderr)
        if "b" not in dbfile.mode:
            raise ValueError("must open file in binary mode")
        # Windows Python 2 file objects can't handle writes >= 4GiB. Objects returned
        # by io.open() work around this issue, see https://bugs.python.org/issue9611
        if not isinstance(dbfile, io.BufferedIOBase) and self._table_bytes >= 1 << 32:
            raise ValueError("must open file with io.open if size >= 4GiB")
        dbfile.truncate(dbfile.tell() + self.HEADER_LEN + self._table_bytes)
        dbfile.write(self._header())
        dbfile.write(self._data)

    @classmethod
    def fromfile(cls, dbfile, mmap_access = mmap.ACCESS_READ, preload = True):
        """Load the address set from a file

        :param dbfile: an open file object from which the set is loaded;
                       it will be closed by AddressSet when no longer needed
        :type dbfile: io.FileIO or file
        :param mmap_access: mmap.ACCESS_READ, .ACCESS_WRITE, or .ACCESS_COPY
        :type mmap_access: int
        :param preload: True to preload the entire address set, False to load on demand
        :type preload: bool
        """
        if "b" not in dbfile.mode:
            raise ValueError("must open file in binary mode")
        header_pos = dbfile.tell()
        if header_pos % mmap.ALLOCATIONGRANULARITY != 0:
            raise ValueError("header position in file must be a multiple of {}".format(mmap.ALLOCATIONGRANULARITY))
        #
        # Read in the header safely (ast.literal_eval() is safe for untrusted data)
        header = dbfile.read(cls.HEADER_LEN)
        if not header.startswith(cls.MAGIC):
            raise ValueError("unrecognized file format (invalid magic)")
        magic_len  = len(cls.MAGIC)
        config_end = header.find(b"\0", magic_len, cls.HEADER_LEN)
        assert config_end > 0
        config = ast.literal_eval(header[magic_len:config_end])
        if config["version"] != cls.VERSION:
            raise ValueError("can't load address database version {} (only supports {})"
                             .format(config["version"], cls.VERSION))
        #
        # Create an AddressSet object and replace its attributes
        self = cls(1)  # (size is irrelevant since it's getting replaced)
        cls._remove_nonheader_attribs(self.__dict__)
        for attr in self.__dict__.keys():  # only load expected attributes from untrusted data
            self.__dict__[attr] = config[attr]
        self._mmap_access = mmap_access
        #
        # The hash table is memory-mapped directly from the file instead of being loaded
        self._data = mmap.mmap(dbfile.fileno(), self._table_bytes, access=mmap_access,
                                offset= header_pos + cls.HEADER_LEN)
        if mmap_access == mmap.ACCESS_WRITE:
            dbfile.seek(header_pos)  # prepare for writing an updated header in close()
        else:
            dbfile.close()
        self._dbfile = dbfile
        #
        # Most of the time it makes sense to load the file serially instead of letting
        # the OS load each page as it's touched in random order, especially with HDDs;
        # reading a byte from each page is sufficient (CPython doesn't optimize this away)
        if preload:
            for i in xrange(self._table_bytes // mmap.PAGESIZE):
                self._data[i * mmap.PAGESIZE]
        #
        return self

    def close(self, flush = True):
        if self._dbfile:                 # if present, self._data is an mmap
            if not self._dbfile.closed:  # if not closed, the mmap was opened in write/update mode
                self._dbfile.write(self._header())  # update the header
                self._dbfile.close()
                if flush:
                    self._data.flush()
            self._data.close()
            self._dbfile = None
        elif isinstance(self._data, bytearray) and self._data:
            self._data = bytearray()
        if flush:
            gc.collect()

    def __del__(self):
        if hasattr(self, "_dbfile"):
            self.close(flush=False)


# Decodes a Bitcoin-style variable precision integer and
# returns a tuple containing its value and incremented offset
def varint(data, offset):
    b = ord(data[offset])
    if b <= 252:
        return b, offset + 1
    if b == 253:
        return struct.unpack_from("<H", data, offset + 1)[0], offset + 3
    if b == 254:
        return struct.unpack_from("<I", data, offset + 1)[0], offset + 5
    if b == 255:
        return struct.unpack_from("<Q", data, offset + 1)[0], offset + 9
    assert False


def create_address_db(dbfilename, blockdir, update = False, progress_bar = True):
    """Creates an AddressSet database and saves it to a file

    :param dbfilename: the file name where the database is saved (overwriting it)
    :type dbfilename: str
    :param blockdir: the data directory where the Bitcoin block files reside
    :type blockdir: str
    :param update: if True, the existing database file is updated from new txs
    :type update: bool
    :param progress_bar: True to enable the progress bar
    :type progress_bar: bool
    """

    for filename in glob.iglob(path.join(blockdir, "blk*.dat")):
        if path.isfile(filename): break
    else:
        raise ValueError("no block files exist in blocks directory '{}'".format(blockdir))

    if update:
        print("Loading address database ...")
        address_set   = AddressSet.fromfile(open(dbfilename, "r+b"), mmap_access=mmap.ACCESS_WRITE)
        first_filenum = address_set.last_filenum
        print()
    else:
        first_filenum = 0

    filename = "blk{:05}.dat".format(first_filenum)
    if not path.isfile(path.join(blockdir, filename)):
        raise ValueError("first block file '{}' doesn't exist in blocks directory '{}'".format(filename, blockdir))

    if not update:
        # Open the file early to make sure we can, but don't overwrite it yet
        # (see AddressSet.tofile() for why io.open() instead of open() is used)
        try:
            dbfile = io.open(dbfilename, "r+b")
        except IOError:
            dbfile = io.open(dbfilename, "wb")
        # With the default bytes_per_addr and max_load, this allocates
        # about 4 GiB which is room for a little over 400 million addresses
        address_set = AddressSet(1 << 29)

    if progress_bar:
        try:
            import progressbar
        except ImportError:
            progress_bar = False

    if progress_bar:
        print("Parsing block files ...")
        for filenum in itertools.count(first_filenum):
            filename = path.join(blockdir, "blk{:05}.dat".format(filenum))
            if not path.isfile(filename):
                break
        progress_label = progressbar.FormatLabel(" {:11,} addrs. %(elapsed)s, ".format(len(address_set)))
        progress_bar = progressbar.ProgressBar(maxval=filenum-first_filenum, widgets=[
            progressbar.SimpleProgress(), " ",
            progressbar.Bar(left="[", fill="-", right="]"),
            progress_label,
            progressbar.ETA()
        ])
        progress_bar.start()
    else:
        print("Block file   Address count")
        print("------------ -------------")
        # e.g. blk00943.dat   255,212,706

    for filenum in itertools.count(first_filenum):
        filename = path.join(blockdir, "blk{:05}.dat".format(filenum))
        if not path.isfile(filename):
            break
        address_set.last_filenum = filenum

        with open(filename, "rb") as blockfile:
            if not progress_bar:
                print(path.basename(filename), end=" ")

            header = blockfile.read(8)  # read in the magic and remaining (after these 8 bytes) block length
            while len(header) == 8 and header[4:] != b"\0\0\0\0":
                assert header[:4] == b"\xf9\xbe\xb4\xd9"                        # magic

                block = blockfile.read(struct.unpack_from("<I", header, 4)[0])  # read in the rest of the block
                tx_count, offset = varint(block, 80)                            # skips 80 bytes of header
                for tx_num in xrange(tx_count):
                    offset += 4                                                 # skips 4-byte tx version
                    is_bip144 = block[offset] == b"\0"                          # bip-144 marker
                    if is_bip144:
                        offset += 2                                             # skips 1-byte marker & 1-byte flag
                    txin_count, offset = varint(block, offset)
                    for txin_num in xrange(txin_count):
                        sigscript_len, offset = varint(block, offset + 36)      # skips 32-byte tx id & 4-byte tx index
                        offset += sigscript_len + 4                             # skips sequence number & sigscript
                    txout_count, offset = varint(block, offset)
                    for txout_num in xrange(txout_count):
                        pkscript_len, offset = varint(block, offset + 8)        # skips 8-byte satoshi count

                        # If this is a P2PKH script (OP_DUP OP_HASH160 PUSH(20) <20 address bytes> OP_EQUALVERIFY OP_CHECKSIG)
                        if pkscript_len == 25 and block[offset:offset+3] == b"\x76\xa9\x14" and block[offset+23:offset+25] == b"\x88\xac":
                            # Add the discovered address to the address set
                            address_set.add(block[offset+3:offset+23])

                        elif pkscript_len == 67 and block[offset] == b"\x41" and block[offset + 66] == b"\xac":
                            address_set.add(bin_hash160(block[offset+1:offset+66]))

                        elif pkscript_len == 35 and block[offset] == b"\x21" and block[offset+34] == b"\xac":
                            address_set.add(bin_hash160(decompress(block[offset+1:offset+34])))

                        offset += pkscript_len                                  # advances past the pubkey script
                    if is_bip144:
                        for txin_num in xrange(txin_count):
                            stackitem_count, offset = varint(block, offset)
                            for stackitem_num in xrange(stackitem_count):
                                stackitem_len, offset = varint(block, offset)
                                offset += stackitem_len                         # skips this stack item
                    offset += 4                                                 # skips the 4-byte locktime
                header = blockfile.read(8)  # read in the next magic and remaining block length

        if progress_bar:
            progress_label.format = " {:11,} addrs. %(elapsed)s, ".format(len(address_set))  # updates address count
            nextval = progress_bar.currval + 1
            if nextval > progress_bar.maxval:  # can happen if the bitcoin client is left running
                progress_bar.maxval = nextval
            progress_bar.update(nextval)
        else:
            print("{:13,}".format(len(address_set)))

    if progress_bar:
        progress_bar.widgets.pop()  # remove the ETA
        progress_bar.finish()
    if update:
        print("\nSaving changes to address database ...")
        address_set.close()
    else:
        print("\nSaving address database ...")
        dbfile.truncate(0)
        address_set.tofile(dbfile)
        dbfile.close()

    print("\nDone.")
