Merge branch 'feature-unzip' into 'master'

Feature: Unzip and decompression

Unzip downloaded file, if its name ends with `.zip`, and has a single file or a single iso file in it. ungz/unbz2 if url ends with `bz2` or `gz`.

disk.py

@@ -3,6 +3,11 @@ import os
 import subprocess
 import re
 import logging
+from shutil import move, copyfileobj
+from zipfile import ZipFile, is_zipfile
+from zlib import decompressobj, MAX_WBITS
+from bz2 import BZ2Decompressor
+
 import requests
 
 logger = logging.getLogger(__name__)
@@ -105,7 +110,7 @@ class Disk(object):
         # Call subprocess
         subprocess.check_output(cmdline)
 
-    def download(self, task, url, parent_id=None):
+    def download(self, task, url, parent_id=None):  # noqa
         ''' Download image from url. '''
         disk_path = self.get_path()
         logger.info("Downloading image from %s to %s", url, disk_path)
@@ -115,34 +120,73 @@ class Disk(object):
                 pass
             if parent_id is None:
                 parent_id = task.request.id
-            percent_size = float(r.headers['content-length']) / 100
-            percent = 0
-            actual_size = 0
             chunk_size = 256 * 1024
+            ext = url.split('.')[-1].lower()
+            if ext == 'gz':
+                decompressor = decompressobj(16 + MAX_WBITS)
+                # undocumented zlib feature http://stackoverflow.com/a/2424549
+            elif ext == 'bz2':
+                decompressor = BZ2Decompressor()
+            clen = max(int(r.headers.get('content-length', 700000000)), 1)
+            percent = 0
             try:
                 with open(disk_path, 'wb') as f:
                     for chunk in r.iter_content(chunk_size=chunk_size):
                         if task.is_aborted():
                             raise AbortException()
-                        if chunk:
+                        if ext in ('gz', 'bz'):
+                            chunk = decompressor.decompress(chunk)
                         f.write(chunk)
-                            actual_size += chunk_size
-                            if actual_size > (percent_size * percent):
-                                percent += 1
+                        actsize = f.tell()
+                        new_percent = min(100, round(actsize * 100.0 / clen))
+                        if new_percent > percent:
+                            percent = new_percent
                             task.update_state(
                                 task_id=parent_id,
                                 state=task.AsyncResult(parent_id).state,
-                                    meta={'size': actual_size,
-                                          'percent': percent})
+                                meta={'size': actsize, 'percent': percent})
+                    if ext == 'gz':
+                        f.write(decompressor.flush())
                     f.flush()
                 self.size = os.path.getsize(disk_path)
                 logger.debug("Download finished %s (%s bytes)",
                              self.name, self.size)
             except AbortException:
                 # Cleanup file:
-                os.unlink(self.get_path())
+                os.unlink(disk_path)
                 logger.info("Download %s aborted %s removed.",
                             url, disk_path)
+            except:
+                os.unlink(disk_path)
+                logger.error("Download %s failed, %s removed.",
+                             url, disk_path)
+                raise
+            else:
+                if ext == 'zip' and is_zipfile(disk_path):
+                    task.update_state(
+                        task_id=parent_id,
+                        state=task.AsyncResult(parent_id).state,
+                        meta={'size': actsize, 'extracting': 'zip',
+                              'percent': 99})
+                    self.extract_iso_from_zip(disk_path)
+
+    def extract_iso_from_zip(self, disk_path):
+        with ZipFile(disk_path, 'r') as z:
+            isos = z.namelist()
+            if len(isos) != 1:
+                isos = [i for i in isos
+                        if i.lower().endswith('.iso')]
+            if len(isos) == 1:
+                logger.info('Unzipping %s started.', disk_path)
+                f = open(disk_path + '~', 'wb')
+                zf = z.open(isos[0])
+                with zf, f:
+                    copyfileobj(zf, f)
+                    f.flush()
+                move(disk_path + '~', disk_path)
+            else:
+                logger.info("Extracting %s failed, keeping original.",
+                            disk_path)
 
     def snapshot(self):
         ''' Creating qcow2 snapshot with base image.