clean up imports; do not shell-out for tar/bzip2/bunzip2. Instead, use tarfile module

blockstack/lib/fast_sync.py

@@ -21,35 +21,20 @@
     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 tarfile
 
 import virtualchain
 from virtualchain.lib.ecdsalib import sign_digest, verify_digest
 
-import blockstack_client
-
 log = virtualchain.get_logger("blockstack-server")
 
-import nameset as blockstack_state_engine
 import nameset.virtualchain_hooks as virtualchain_hooks
 
 import config
@@ -194,6 +179,66 @@ def fast_sync_sign_snapshot( snapshot_path, private_key, first=False ):
     return True
 
 
+def fast_sync_snapshot_compress( snapshot_dir, export_path ):
+    """
+    Given the path to a directory, compress it and export it to the
+    given path.
+
+    Return {'status': True} on success
+    Return {'error': ...} on failure
+    """
+
+    snapshot_dir = os.path.abspath(snapshot_dir)
+    export_path = os.path.abspath(export_path)
+    if os.path.exists(export_path):
+        return {'error': 'Snapshot path exists: {}'.format(export_path)}
+
+    old_dir = os.getcwd()
+    
+    count_ref = [0]
+
+    def print_progress(tarinfo):
+        count_ref[0] += 1
+        if count_ref[0] % 100 == 0:
+            log.debug("{} files...".format(count_ref[0]))
+
+        return tarinfo
+
+    try:
+        os.chdir(snapshot_dir)
+        with tarfile.TarFile.bz2open(export_path, "w") as f:
+            f.add(".", filter=print_progress)
+
+    except:
+        os.chdir(old_dir)
+        raise
+    
+    finally:
+        os.chdir(old_dir)
+
+    return {'status': True}
+
+
+def fast_sync_snapshot_decompress( snapshot_path, output_dir ):
+    """
+    Given the path to a snapshot file, decompress it and 
+    write its contents to the given output directory
+
+    Return {'status': True} on success
+    Return {'error': ...} on failure
+    """
+    if not tarfile.is_tarfile(snapshot_path):
+        return {'error': 'Not a tarfile-compatible archive: {}'.format(snapshot_path)}
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    with tarfile.TarFile.bz2open(snapshot_path, 'r') as f:
+        tarfile.TarFile.extractall(f, path=output_dir)
+
+    return {'status': True}
+
+
 def fast_sync_snapshot( export_path, private_key, block_number ):
     """
     Export all the local state for fast-sync.
@@ -263,7 +308,7 @@ def fast_sync_snapshot( export_path, private_key, block_number ):
     _zonefile_copy_progress = _zonefile_copy_progress_outer()
 
     # make sure we have the apppriate tools
-    tools = ['tar', 'bzip2', 'mv', 'sqlite3']
+    tools = ['sqlite3']
     for tool in tools:
         rc = os.system("which {} > /dev/null".format(tool))
         if rc != 0:
@@ -339,16 +384,15 @@ def fast_sync_snapshot( export_path, private_key, block_number ):
 
     # 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))
+    res = fast_sync_snapshot_compress(tmpdir, export_path)
+    if 'error' in res:
+        log.error("Faield to compress {} to {}: {}".format(tmpdir, export_path, res['error']))
         _cleanup(tmpdir)
         return False
 
     log.debug("Wrote {} bytes".format(os.stat(export_path).st_size))
 
+    # sign
     rc = fast_sync_sign_snapshot( export_path, private_key, first=True )
     if not rc:
         log.error("Failed to sign snapshot {}".format(export_path))
@@ -470,14 +514,6 @@ def fast_sync_import( working_dir, import_url, public_keys=config.FAST_SYNC_PUBL
     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()
 
@@ -530,8 +566,8 @@ def fast_sync_import( working_dir, import_url, public_keys=config.FAST_SYNC_PUBL
                     
                     log.debug("Public key {} matches {} ({})".format(next_pubkey, sigb64, hash_hex))
                     signatures.remove(sigb64)
-
-                elif os.environ.get("BLOCKSTACK_TEST") == "1":
+                
+                else:
                     log.debug("Public key {} does NOT match {} ({})".format(next_pubkey, sigb64, hash_hex))
 
         # enough signatures?
@@ -541,11 +577,9 @@ def fast_sync_import( working_dir, import_url, public_keys=config.FAST_SYNC_PUBL
 
     # 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))
+    res = fast_sync_snapshot_decompress(import_path, working_dir)
+    if 'error' in res:
+        log.error("Failed to decompress {} to {}: {}".format(import_path, working_dir, res['error']))
         return False
 
     # restore from backup