stacks-puppet-node/blockstack/lib/fast_sync.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
    Blockstack
    ~~~~~
    copyright: (c) 2014-2015 by Halfmoon Labs, Inc.
    copyright: (c) 2016-2017 by Blockstack.org

    This file is part of Blockstack

    Blockstack 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 3 of the License, or
    (at your option) any later version.

    Blockstack 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 Blockstack. If not, see <http://www.gnu.org/licenses/>.
"""

import logging
import os
import sys
import json
import datetime
import traceback
import time
import math
import random
import shutil
import tempfile
import binascii
import copy
import threading
import errno
import base64
import keylib
import subprocess
import urllib
import hashlib

import virtualchain
import blockstack_client

log = virtualchain.get_logger("blockstack-server")

import pybitcoin

import nameset as blockstack_state_engine
import nameset.virtualchain_hooks as virtualchain_hooks

import config

from .b40 import *
from .config import *
from .scripts import *
from .hashing import *
from .storage import *

from .nameset import *
from .operations import *

def snapshot_peek_number( fd, off ):
    """
    Read the last 8 bytes of fd
    and interpret it as an int.
    """
    # read number of 8 bytes 
    fd.seek( off - 8, os.SEEK_SET )
    value_hex = fd.read(8)
    if len(value_hex) != 8:
        return None
    try:
        value = int(value_hex, 16)
    except ValueError:
        return None

    return value


def snapshot_peek_sigb64( fd, off, bytelen ):
    """
    Read the last :bytelen bytes of
    fd and interpret it as a base64-encoded
    string
    """
    fd.seek( off - bytelen, os.SEEK_SET )
    sigb64 = fd.read(bytelen)
    if len(sigb64) != bytelen:
        return None

    try:
        base64.b64decode(sigb64)
    except:
        return None

    return sigb64


def fast_sync_sign_snapshot( snapshot_path, private_key, first=False ):
    """
    Append a signature to the end of a snapshot path
    with the given private key.

    If first is True, then don't expect the signature trailer.

    Return True on success
    Return False on error
    """
   
    if not os.path.exists(snapshot_path):
        log.error("No such file or directory: {}".format(snapshot_path))
        return False

    file_size = 0
    payload_size = 0
    write_offset = 0
    try:
        sb = os.stat(snapshot_path)
        file_size = sb.st_size
        assert file_size > 8
    except Exception as e:
        log.exception(e)
        return False
    
    num_sigs = 0
    snapshot_hash = None
    with open(snapshot_path, 'r+') as f:

        if not first:
            info = fast_sync_inspect(f)
            if 'error' in info:
                log.error("Failed to inspect {}: {}".format(snapshot_path, info['error']))
                return False

            num_sigs = len(info['signatures'])
            write_offset = info['sig_append_offset']
            payload_size = info['payload_size']
 
        else:
            # no one has signed yet.
            write_offset = file_size
            num_sigs = 0
            payload_size = file_size

        # hash the file and sign the (bin-encoded) hash
        privkey_hex = keylib.ECPrivateKey(private_key).to_hex()
        hash_hex = blockstack_client.storage.get_file_hash( f, hashlib.sha256, fd_len=payload_size )
        sigb64 = blockstack_client.keys.sign_digest( hash_hex, privkey_hex, hashfunc=hashlib.sha256 )
      
        if os.environ.get("BLOCKSTACK_TEST") == "1":
            log.debug("Signed {} with {} to make {}".format(hash_hex, keylib.ECPrivateKey(private_key).public_key().to_hex(), sigb64))

        # append
        f.seek(write_offset, os.SEEK_SET)
        f.write(sigb64)
        f.write('{:08x}'.format(len(sigb64)))

        # append number of signatures
        num_sigs += 1
        f.write('{:08x}'.format(num_sigs))
    
        f.flush()
        os.fsync(f.fileno())

    return True


def fast_sync_snapshot( export_path, private_key, block_number ):
    """
    Export all the local state for fast-sync.
    If block_number is given, then the name database
    at that particular block number will be taken.

    The exported tarball will be signed with the given private key,
    and the signature will be appended to the end of the file.

    Return True if we succeed
    Return False if not
    """

    db_paths = None
    found = True
    tmpdir = None
    namedb_path = None

    def _cleanup(path):
        try:
            # shutil.rmtree(path)
            print 'rm -rf {}'.format(path)
        except Exception, e:
            log.exception(e)
            log.error("Failed to clear directory {}".format(path))

    
    def _log_backup(path):
        sb = None
        try:
            sb = os.stat(path)
        except Exception, e:
            log.exception(e)
            log.error("Failed to stat {}".format(path))
            return False

        log.debug("Back up {} ({} bytes)".format(path, sb.st_size))


    def _copy_paths(src_paths, dest_dir):
        for db_path in src_paths:
            dest_path = os.path.join(dest_dir, os.path.basename(db_path))
            try:
                _log_backup(db_path)
                shutil.copy(db_path, dest_path)
            except Exception, e:
                log.exception(e)
                log.error("Failed to copy {} to {}".format(db_path, dest_path))
                return False
        
        return True


    # ugly hack to work around the lack of a `nonlocal` keyword in Python 2.x
    def _zonefile_copy_progress_outer():
        def inner(src, names):
            for name in names:
                if name == 'zonefile.txt':
                    inner.zonefile_count += 1
                    if inner.zonefile_count % 100 == 0:
                        log.debug("{} zone files copied".format(inner.zonefile_count))
        
            return []

        inner.zonefile_count = 0
        return inner

    _zonefile_copy_progress = _zonefile_copy_progress_outer()

    # make sure we have the apppriate tools
    tools = ['tar', 'bzip2', 'mv', 'sqlite3']
    for tool in tools:
        rc = os.system("which {} > /dev/null".format(tool))
        if rc != 0:
            log.error("'{}' command not found".format(tool))
            return False

    working_dir = virtualchain.get_working_dir() 
    if not os.path.exists(working_dir):
        log.error("No such directory {}".format(working_dir))
        return False

    if block_number is None:
        # last backup
        all_blocks = BlockstackDB.get_backup_blocks( virtualchain_hooks )
        if len(all_blocks) == 0:
            log.error("No backups available")
            return False

        block_number = max(all_blocks)

    log.debug("Snapshot from block {}".format(block_number))

    # use a backup database 
    db_paths = BlockstackDB.get_backup_paths( block_number, virtualchain_hooks )

    for p in db_paths:
        if not os.path.exists(p):
            log.error("Missing file: '%s'" % p)
            found = False

    if not found:
        return False

    try:
        tmpdir = tempfile.mkdtemp(prefix='.blockstack-export-')
    except Exception, e:
        log.exception(e)
        return False

    # copying from backups 
    backups_path = os.path.join(tmpdir, "backups")
    try:
        os.makedirs(backups_path)
    except Exception, e:
        log.exception(e)
        log.error("Failed to make directory {}".format(backups_path))
        _cleanup(tmpdir)
        return False

    rc = _copy_paths(db_paths, backups_path)
    if not rc:
        _cleanup(tmpdir)
        return False

    # copy over atlasdb
    atlasdb_path = os.path.join(working_dir, "atlas.db")
    dest_path = os.path.join(tmpdir, "atlas.db")
    _log_backup(atlasdb_path)
    rc = sqlite3_backup(atlasdb_path, dest_path)
    if not rc:
        _cleanup(tmpdir)
        return False

    # copy over zone files
    zonefiles_path = os.path.join(working_dir, "zonefiles")
    dest_path = os.path.join(tmpdir, "zonefiles")
    try:
        shutil.copytree(zonefiles_path, dest_path, ignore=_zonefile_copy_progress)
    except Exception, e:
        log.exception(e)
        log.error('Failed to copy {} to {}'.format(zonefiles_path, dest_path))
        return False

    # compress
    export_path = os.path.abspath(export_path)
    cmd = "cd '{}' && tar cf 'snapshot.tar' * && bzip2 'snapshot.tar' && mv 'snapshot.tar.bz2' '{}'".format(tmpdir, export_path)
    log.debug("Compressing: {}".format(cmd))
    rc = os.system(cmd)
    if rc != 0:
        log.error("Failed to compress {}. Exit code {}. Command: \"{}\"".format(tmpdir, rc, cmd))
        _cleanup(tmpdir)
        return False

    log.debug("Wrote {} bytes".format(os.stat(export_path).st_size))

    rc = fast_sync_sign_snapshot( export_path, private_key, first=True )
    if not rc:
        log.error("Failed to sign snapshot {}".format(export_path))
        return False

    return True


def fast_sync_fetch( import_url ):
    """
    Get the data for an import snapshot.
    Store it to a temporary path
    Return the path on success
    Return None on error
    """
    try:
        fd, tmppath = tempfile.mkstemp(prefix='.blockstack-fast-sync-')
    except Exception, e:
        log.exception(e)
        return None
    
    try:
        path, headers = urllib.urlretrieve(import_url, tmppath)
    except Exception, e:
        os.close(fd)
        log.exception(e)
        return None
    
    os.close(fd)
    return tmppath


def fast_sync_inspect( fd ):
    """
    Inspect a snapshot, given its file descriptor.
    Get the signatures and payload size
    Return {'status': True, 
            'signatures': signatures,
            'payload_size': payload size,
            'sig_append_offset': offset} on success
    Return {'error': ...} on error
    """
    sb = os.fstat(fd.fileno())
    ptr = sb.st_size
    if ptr < 8:
        log.debug("fd is {} bytes".format(ptr))
        return {'error': 'File is too small to be a snapshot'}

    signatures = []
    sig_append_offset = 0
    
    fd.seek(0, os.SEEK_SET)

    # read number of signatures
    num_signatures = snapshot_peek_number(fd, ptr)
    if num_signatures is None or num_signatures > 256:
        log.error("Unparseable num_signatures field")
        return {'error': 'Unparseable num_signatures'}

    # consumed
    ptr -= 8

    # future signatures get written here
    sig_append_offset = ptr

    # read signatures
    for i in xrange(0, num_signatures):
        sigb64_len = snapshot_peek_number(fd, ptr)
        if sigb64_len is None or sigb64_len > 100:
            log.error("Unparseable signature length field")
            return {'error': 'Unparseable signature length'}

        # consumed length
        ptr -= 8

        sigb64 = snapshot_peek_sigb64(fd, ptr, sigb64_len)
        if sigb64 is None:
            log.error("Unparseable signature")
            return {'error': 'Unparseable signature'}

        # consumed signature
        ptr -= len(sigb64)

        signatures.append( sigb64 )

    return {'status': True, 'signatures': signatures, 'payload_size': ptr, 'sig_append_offset': sig_append_offset}


def fast_sync_import( working_dir, import_url, public_keys=config.FAST_SYNC_PUBLIC_KEYS, num_required=len(config.FAST_SYNC_PUBLIC_KEYS) ):
    """
    Fast sync import.
    Verify the given fast-sync file from @import_path using @public_key, and then 
    uncompress it into @working_dir.

    Verify that at least `num_required` public keys in `public_keys` signed.
    NOTE: `public_keys` needs to be in the same order as the private keys that signed.
    """

    # make sure we have the apppriate tools
    tools = ['tar', 'bzip2', 'mv']
    for tool in tools:
        rc = os.system("which {} > /dev/null".format(tool))
        if rc != 0:
            log.error("'{}' command not found".format(tool))
            return False

    if working_dir is None:
        working_dir = virtualchain.get_working_dir()

    if not os.path.exists(working_dir):
        log.error("No such directory {}".format(working_dir))
        return False

    # go get it 
    import_path = fast_sync_fetch(import_url)
    if import_path is None:
        log.error("Failed to fetch {}".format(import_url))
        return False

    # format: <signed bz2 payload> <sigb64> <sigb64 length (8 bytes hex)> ... <num signatures>
    file_size = 0
    try:
        sb = os.stat(import_path)
        file_size = sb.st_size
    except Exception as e:
        log.exception(e)
        return False

    num_signatures = 0
    ptr = file_size
    signatures = []

    with open(import_path, 'r') as f:
        info = fast_sync_inspect( f )
        if 'error' in info:
            log.error("Failed to inspect snapshot {}: {}".format(import_path, info['error']))
            return False

        signatures = info['signatures']
        ptr = info['payload_size']

        # get the hash of the file 
        hash_hex = blockstack_client.storage.get_file_hash(f, hashlib.sha256, fd_len=ptr)
        
        # validate signatures over the hash
        log.debug("Verify {} bytes".format(ptr))
        key_idx = 0
        num_match = 0
        for next_pubkey in public_keys:
            for sigb64 in signatures:
                valid = blockstack_client.keys.verify_digest( hash_hex, keylib.ECPublicKey(next_pubkey).to_hex(), sigb64, hashfunc=hashlib.sha256 ) 
                if valid:
                    num_match += 1
                    if num_match >= num_required:
                        break
                    
                    log.debug("Public key {} matches {} ({})".format(next_pubkey, sigb64, hash_hex))
                    signatures.remove(sigb64)

                elif os.environ.get("BLOCKSTACK_TEST") == "1":
                    log.debug("Public key {} does NOT match {} ({})".format(next_pubkey, sigb64, hash_hex))

        # enough signatures?
        if num_match < num_required:
            log.error("Not enough signatures match (required {}, found {})".format(num_required, num_match))
            return False

    # decompress
    import_path = os.path.abspath(import_path)
    cmd = "cd '{}' && tar xf '{}'".format(working_dir, import_path)
    log.debug(cmd)
    rc = os.system(cmd)
    if rc != 0:
        log.error("Failed to decompress. Exit code {}. Command: {}".format(rc, cmd))
        return False

    # restore from backup
    rc = blockstack_backup_restore(working_dir, None)
    if not rc:
        log.error("Failed to instantiate blockstack name database")
        return False

    # success!
    log.debug("Restored to {}".format(working_dir))
    return True


def blockstack_backup_restore( working_dir, block_number ):
    """
    Restore the database from a backup in the backups/ directory.
    If block_number is None, then use the latest backup.
    Return True on success
    Return False on failure
    """

    # TODO: this is pretty shady...
    old_working_dir = os.environ.get('VIRTUALCHAIN_WORKING_DIR', None)
    os.environ['VIRTUALCHAIN_WORKING_DIR'] = working_dir

    if block_number is None:
        all_blocks = BlockstackDB.get_backup_blocks( virtualchain_hooks )
        if len(all_blocks) == 0:
            log.error("No backups available")
    
            # TODO: this is pretty shady...
            if old_working_dir:
                os.environ['VIRTUALCHAIN_WORKING_DIR'] = old_working_dir

            return False

        block_number = max(all_blocks)

    found = True
    backup_paths = BlockstackDB.get_backup_paths( block_number, virtualchain_hooks )
    for p in backup_paths:
        if not os.path.exists(p):
            log.error("Missing backup file: '%s'" % p)
            found = False

    if not found:

        # TODO: this is pretty shady...
        if old_working_dir:
            os.environ['VIRTUALCHAIN_WORKING_DIR'] = old_working_dir

        return False 

    rc = BlockstackDB.backup_restore( block_number, virtualchain_hooks )
    if not rc:
        log.error("Failed to restore backup")

        # TODO: this is pretty shady...
        if old_working_dir:
            os.environ['VIRTUALCHAIN_WORKING_DIR'] = old_working_dir
            
        return False

    log.debug("Restored backup from {}".format(block_number))

    # TODO: this is pretty shady...
    if old_working_dir:
        os.environ['VIRTUALCHAIN_WORKING_DIR'] = old_working_dir

    return True