uhd: re-wrote uhd_images_downloader (OO, new options, hopefully much more robust and user-friendly)

1 files changed, 316 insertions, 119 deletions

diff --git a/host/utils/uhd_images_downloader.py.in b/host/utils/uhd_images_downloader.py.in
index bb082190c..697bd4e16 100644
--- a/host/utils/uhd_images_downloader.py.in
+++ b/host/utils/uhd_images_downloader.py.in
@@ -16,144 +16,341 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-import atexit
-import hashlib
+import sys, os, string, tempfile, math
+import traceback
+import shutil, hashlib, urllib2, zipfile
+
 from optparse import OptionParser
-import os
-import os.path
-import shutil
-import string
-import sys
-import tempfile
-import urllib2
-import zipfile
+
+_DEFAULT_BUFFER_SIZE      = 8192
+_BASE_DIR_STRUCTURE_PARTS = ["share", "uhd", "images"]
+_BASE_DIR_STRUCTURE       = os.path.join(_BASE_DIR_STRUCTURE_PARTS)
+_DEFAULT_INSTALL_PATH     = os.path.join("@CMAKE_INSTALL_PREFIX@", *_BASE_DIR_STRUCTURE)
+_AUTOGEN_IMAGES_SOURCE    = "@UHD_IMAGES_DOWNLOAD_SRC@"
+_AUTOGEN_IMAGES_CHECKSUM  = "@UHD_IMAGES_MD5SUM@"
+_IMAGES_CHECKSUM_TYPE     = "md5"
+_CONTACT                  = "support@ettus.com"
 
 def md5Checksum(filePath):
-    with open(filePath, 'rb') as fh:
-        m = hashlib.md5()
-        while True:
-            data = fh.read(8192)
-            if not data:
-                break
-            m.update(data)
-        return m.hexdigest()
+    try:
+        with open(filePath, 'rb') as fh:
+            m = hashlib.md5()
+            while True:
+                data = fh.read(_DEFAULT_BUFFER_SIZE)
+                if not data:
+                    break
+                m.update(data)
+            return m.hexdigest()
+    except Exception, e:
+        print "Failed to calculated MD5 sum of: %s (%s)" % (filePath, e)
+        raise e
 
-class temp_dir():
+_checksum_fns = {
+    'md5': md5Checksum
+}
 
+class temporary_directory():
     def __enter__(self):
-        self.name = tempfile.mkdtemp()
-        return self.name
+        try:
+            self.name = tempfile.mkdtemp()
+            return self.name
+        except Exception, e:
+            print "Failed to create a temporary directory (%s)" % (e)
+            raise e
+    
+    # Can return 'True' to suppress incoming exception
     def __exit__(self, type, value, traceback):
         try:
             shutil.rmtree(self.name)
-        except OSError,e:
-            #Utility should have already detected this, but this is for safety
-            print str(e)
-            raise Exception("Could not install images! Make sure you have write permissions.")
-
-if __name__ == "__main__":
-
-    print
-    if os.environ.get("UHD_IMAGES_DIR") != None and os.environ.get("UHD_IMAGES_DIR") != "":
-        default_images_dir = os.environ.get("UHD_IMAGES_DIR")
-        print "UHD_IMAGES_DIR environment variable is set. Default install location: %s" % default_images_dir
-    else:
-        default_images_dir = "@CMAKE_INSTALL_PREFIX@/share/uhd/images"
+        except Exception, e:
+            print "Could not delete temporary directory: %s (%s)" % (self.name, e)
 
-    #Command line options
-    parser = OptionParser()
-    parser.add_option("--install-location", type="string", default=default_images_dir, help="Set custom install location for images")
-    parser.add_option("--buffer-size", type="int", default=8192, help="Set download buffer size, [default=%default]",)
-    (options, args) = parser.parse_args()
+class uhd_images_downloader():
+    def __init__(self):
+        pass
     
-    #Configuring image download info
-    images_src = "@UHD_IMAGES_DOWNLOAD_SRC@"
-    images_zip_md5sum = "@UHD_IMAGES_MD5SUM@"
-    filename = images_src.split("/")[-1]
-
-    #Use this directory with relative paths
-    current_directory = os.getcwd()
-
-    with temp_dir() as dirname:
-        os.chdir(dirname)
-
-        if os.path.isabs(options.install_location):
-            #Custom absolute path given
-            images_dir = options.install_location
-        else:
-            #Custom relative path given, so construct absolute path
-            images_dir = os.path.abspath(os.path.join(current_directory, options.install_location))
-
-        #Before doing anything, check for write permissions in parent directory
-        parent_directory = os.path.dirname(images_dir)
-        if os.access(parent_directory, os.W_OK):
-            print "Downloading images to: %s" % images_dir
-        else:
-            print "You do not have write permissions at the install location!"
-            sys.exit(1)
-
+    def download(self, images_url, filename, buffer_size=_DEFAULT_BUFFER_SIZE, print_progress=False):
         opener = urllib2.build_opener()
         opener.add_headers = [('User-Agent', 'UHD Images Downloader')]
-        u = opener.open(images_src)
-        f = open(filename, "wb")
+        u = opener.open(images_url)
         meta = u.info()
         filesize = float(meta.getheaders("Content-Length")[0])
         
-        print "Downloading images from: %s" % images_src
+        filesize_dl = 0
         
-        filesize_dl = 0.0
-
-        #Downloading file    
-        while True:
-            buffer = u.read(options.buffer_size)
-            if not buffer:
-                break
-        
-            filesize_dl -= len(buffer)
-            f.write(buffer)
+        with open(filename, "wb") as f:
+            while True:
+                buffer = u.read(buffer_size)
+                if not buffer:
+                    break
+                
+                f.write(buffer)
 
-            status = r"%2.2f MB/%2.2f MB (%3.2f" % (-filesize_dl/1e6, filesize/1e6, (-filesize_dl*100.)/filesize) + r"%)"
-            status += chr(8)*(len(status)+1)
-            print status,
+                filesize_dl += len(buffer)
+                
+                if print_progress:
+                    status = r"%05d kB / %05d kB (%03d%%)" % (int(math.ceil(filesize_dl/1000.)), int(math.ceil(filesize/1000.)), int(math.ceil(filesize_dl*100.)/filesize))
+                    status += chr(8)*(len(status)+1)
+                    print status,
         
-        f.close()
-
-        #Checking md5sum of zip file
-        downloaded_zip_md5sum = md5Checksum(filename)
-        if images_zip_md5sum != downloaded_zip_md5sum:
-            print "\nMD5 checksum does not match!"
-            print "Expected %s, got %s" % (images_zip_md5sum, downloaded_zip_md5sum)
-            print "Images did not install. If problem persists, please contact support@ettus.com."
-            os.remove(filename)
-            os.chdir("/".join(images_dir.split("/")[:-1]))
-            sys.exit(1)
+        if print_progress:
+            print
+        
+        return (filesize, filesize_dl)
+    
+    def check_directories(self, dirs, print_progress=False):
+        if dirs is None or dirs == "":
+            dirs = "."
+        dirs = os.path.abspath(dirs)
+        
+        def _check_part(head, tail=None):
+            if print_progress:
+                print "Checking: %s" % (head)
+            if tail is not None and tail == "":
+                return True
+            if not os.path.exists(head):
+                if print_progress:
+                    print "Does not exist: %s" % (head)
+                return _check_part(*os.path.split(head))
+            if not os.path.isdir(head):
+                if print_progress:
+                    print "Is not a directory: %s" % (head)
+                return (False, head)
+            if not os.access(head, os.W_OK):
+                if print_progress:
+                    print "Write permission denied on: %s" % (head)
+                return (False, head)
+            if print_progress:
+                print "Write permission granted on: %s" % (head)
+            return (True, head)
+        
+        return _check_part(dirs)
+    
+    def validate_checksum(self, checksum_fn, file_path, expecting, print_progress=False):
+        if checksum_fn is None:
+            return (True, "")
+        
+        calculated_checksum = checksum_fn(file_path)
+        
+        if (expecting is not None) and (expecting != "") and calculated_checksum != expecting:
+            return (False, calculated_checksum)
+        
+        return (True, calculated_checksum)
+    
+    def extract_images_archive(self, archive_path, destination=None, print_progress=False):
+        if not os.path.exists(archive_path):
+            if print_progress:
+                print "Path does not exist: %s" % (archive_path)
+            raise Exception("path does not exist: %s" % (archive_path))
+        
+        if print_progress:
+            print "Archive path: %s" % (archive_path)
+        
+        (head, tail) = os.path.split(archive_path)
+        
+        if not os.access(head, os.W_OK):
+            if print_progress:
+                print "Write access denied on: %s" % (head)
+            raise Exception("write access denied on: %s" % (head))
+        
+        (root, ext) = os.path.splitext(tail)
+        temp_dir = os.path.join(head, root)
+        
+        if print_progress:
+            print "Temporary extraction location: %s" % (temp_dir)
+        
+        if os.path.exists(temp_dir):
+            if print_progress:
+                print "Deleting existing location: %s" % (temp_dir)
+            shutil.rmtree(temp_dir)
+        
+        if print_progress:
+            print "Creating directory: %s" % (temp_dir)
+        os.mkdir(temp_dir)
+        
+        if print_progress:
+            print "Extracting archive %s to %s" % (archive_path, temp_dir)
+        
+        images_zip = zipfile.ZipFile(archive_path)
+        images_zip.extractall(temp_dir)
+        images_zip.close()
+        
+        return temp_dir
+    
+    def install_images(self, source, dest, keep=False, print_progress=False):
+        if not os.path.exists(source):
+            if print_progress:
+                print "Source path does not exist: %s" % (source)
+            return
+        
+        if keep:
+            if print_progress:
+                print "Not wiping directory tree (existing files will be overwritten): %s" % (dest)
+        elif os.path.exists(dest):
+            if print_progress:
+                print "Deleting directory tree: %s" % (dest)
+            shutil.rmtree(dest)
+        
+        (head, tail) = os.path.split(source)
+        
+        if print_progress:
+            print "Source install path: %s" % (source)
+        
+        uhd_source = os.path.join(source, tail, *_BASE_DIR_STRUCTURE_PARTS)
+        
+        if print_progress:
+            print "Copying files from: %s" % (uhd_source)
+            print "Copying files to:   %s" % (dest)
+        
+        if keep:
+            # mgrant @ http://stackoverflow.com/questions/12683834/how-to-copy-directory-recursively-in-python-and-overwrite-all
+            def _recursive_overwrite(src, dest, ignore=None):
+                if os.path.isdir(src):
+                    if not os.path.isdir(dest):
+                        os.makedirs(dest)
+                    files = os.listdir(src)
+                    if ignore is not None:
+                        ignored = ignore(src, files)
+                    else:
+                        ignored = set()
+                    for f in files:
+                        if f not in ignored:
+                            _recursive_overwrite(os.path.join(src, f), os.path.join(dest, f), ignore)
+                else:
+                    shutil.copyfile(src, dest)
+            
+            _recursive_overwrite(uhd_source, dest)
         else:
-            temp_path = "tempdir"
-
-            #Extracting contents of zip file
-            if os.path.exists(temp_path):
-                shutil.rmtree(temp_path)
-            os.mkdir(temp_path)
-
-            images_zip = zipfile.ZipFile(filename)
-            images_zip.extractall(temp_path)
+            shutil.copytree(uhd_source, dest)
 
-            #Removing images currently in images_dir
-            if os.path.exists(images_dir):
-                try:
-                    shutil.rmtree(images_dir)
-                except OSError,e:
-                    print str(e)
-                    print "Make sure you have write permissions in the images directory."
-                    sys.exit(1)
-
-            #Copying downloaded images into images_dir
-            shutil.copytree(os.path.join(temp_path, os.path.splitext(filename)[0], 'share', 'uhd', 'images'), images_dir)
-
-            #Removing tempdir and zip file
-            shutil.rmtree(temp_path)
-            images_zip.close()
-            os.remove(filename)
+def main():
+    if os.environ.get("UHD_IMAGES_DIR") != None and os.environ.get("UHD_IMAGES_DIR") != "":
+        default_images_dir = os.environ.get("UHD_IMAGES_DIR")
+        print "UHD_IMAGES_DIR environment variable is set. Default install location: %s" % default_images_dir
+    else:
+        default_images_dir = _DEFAULT_INSTALL_PATH
+    
+    parser = OptionParser()
+    parser.add_option("-i", "--install-location",   type="string",          default=default_images_dir,
+                        help="Set custom install location for images [default=%default]")
+    parser.add_option("--buffer-size",              type="int",             default=_DEFAULT_BUFFER_SIZE,
+                        help="Set download buffer size [default=%default]")
+    parser.add_option("-u", "--url",                type="string",          default=_AUTOGEN_IMAGES_SOURCE,
+                        help="Set images download location [default=%default]")
+    parser.add_option("-c", "--checksum",           type="string",          default=_AUTOGEN_IMAGES_CHECKSUM,
+                        help="Validate images archive against this checksum (blank to skip) [default=%default]")
+    parser.add_option("-t", "--checksum-type",      type="string",          default=_IMAGES_CHECKSUM_TYPE,
+                        help=("Select checksum hash function (options: %s) [default=%%default]" % (",".join(_checksum_fns.keys()))))
+    parser.add_option("-k", "--keep",               action="store_true",    default=False,
+                        help="Do not clear images directory before extracting new files [default=%default]")
+    parser.add_option("-v", "--verbose",            action="store_true",    default=False,
+                        help="Enable verbose output [default=%default]")
+    
+    (options, args) = parser.parse_args()
+    
+    if options.buffer_size <= 0:
+        print "Invalid buffer size: %s" % (options.buffer_size)
+        return 1
+    
+    checksum_fn = None
+    if options.checksum != "":
+        options.checksum_type = options.checksum_type.lower()
+        if not _checksum_fns.has_key(options.checksum_type):
+            print "Not a supported checksum function: %s" % (options.checksum_type)
+            return 1
+        checksum_fn = _checksum_fns[options.checksum_type]
+    
+    url_parts = options.url.split("/")
+    if len(url_parts) <= 1 or url_parts[-1] == "":
+        print "Not a valid URL: %s" % (options.url)
+        return 1
+    images_filename = url_parts[-1]
+    
+    images_dir = os.path.abspath(options.install_location)  # This will use the current working directory if it's not absolute
+    
+    if options.verbose:
+        print "Requested install location: %s" % (options.install_location)
+        print "Images source:              %s" % (options.url)
+        print "Images filename:            %s" % (images_filename)
+        print "Images checksum:            %s (%s)" % (options.checksum, _IMAGES_CHECKSUM_TYPE)
+        print "Final install location:     %s" % (images_dir)
+    else:
+        print "Images destination:      %s" % (images_dir)
+    
+    downloader = uhd_images_downloader()
+    
+    try:
+        (access, last_path) = downloader.check_directories(images_dir, print_progress=options.verbose)
+        if access:
+            with temporary_directory() as temp_dir:
+                if options.verbose:
+                    print "Using temporary directory: %s" % (temp_dir)
+                
+                print "Downloading images from: %s" % options.url
+                
+                temp_images_dest = os.path.join(temp_dir, images_filename)
+                
+                print "Downloading images to:   %s" % (temp_images_dest)
+                
+                (reported_size, downloaded_size) = downloader.download(images_url=options.url, filename=temp_images_dest, buffer_size=options.buffer_size, print_progress=True)
+                
+                if options.verbose:
+                    print "Downloaded %d of %d bytes" % (downloaded_size, reported_size)
+                
+                (checksum_match, calculated_checksum) = downloader.validate_checksum(checksum_fn, temp_images_dest, options.checksum, print_progress=options.verbose)
+                
+                if options.verbose:
+                    print "Calculated checksum: %s" % (calculated_checksum)
+                
+                if checksum_match:
+                    if options.verbose:
+                        if options.checksum == "":
+                            print "Ignoring checksum"
+                        else:
+                            print "Checksum OK"
+                    
+                    try:
+                        extract_path = downloader.extract_images_archive(temp_images_dest, print_progress=options.verbose)
+                        
+                        if options.verbose:
+                            print "Image archive extracted to: %s" % (extract_path)
+                        
+                        downloader.install_images(extract_path, images_dir, options.keep, print_progress=options.verbose)
+                        
+                        print
+                        print "Images successfully installed to: %s" % (images_dir)
+                    except Exception, e:
+                        print "Failed to install image archive: %s" % (e)
+                        print "This is usually a permissions problem."
+                        print "Please check your file system access rights and try again."
+                        
+                        if options.verbose:
+                            traceback.print_exc()
+                        else:
+                            print "You can run this again with the '--verbose' flag to see more information"
+                        print "If the problem persists, please email the output to: %s" % (_CONTACT)
+                else:
+                    print "Checksum of downloaded file is not correct (not installing - see options to override)"
+                    print "Expected:   %s" % (options.checksum)
+                    print "Calculated: %s" % (calculated_checksum)
+                    print "Please try downloading again."
+                    print "If the problem persists, please email the output to: %s" % (_CONTACT)
+        else:
+            print "You do not have sufficient permissions to write to: %s" % (last_path)
+            print "Are you root?"
+    except KeyboardInterrupt:
+        print
+        print "Cancelled at user request"
+    except Exception, e:
+        print "Downloader raised an unhandled exception: %s" % (e)
+        if options.verbose:
+            traceback.print_exc()
+        else:
+            print "You can run this again with the '--verbose' flag to see more information"
+        print "If the problem persists, please email the output to: %s" % (_CONTACT)
+        return 1
+    
+    return 0
 
-            os.chdir(images_dir)
-            print "\n\nImages successfully installed!"
+if __name__ == "__main__":
+    sys.exit(main())