Merge branch 'disk-download-enhancement' into 'master'

Disk Download Enhancement

disk.py

@@ -3,6 +3,7 @@ import os
 import subprocess
 import re
 import logging
+import magic
 from shutil import move, copy, copyfileobj
 from zipfile import ZipFile, is_zipfile
 from zlib import decompressobj, MAX_WBITS
@@ -116,12 +117,30 @@ class Disk(object):
         # Call subprocess
         subprocess.check_output(cmdline)
 
+    def check_valid_image(self):
+        """Check wether the downloaded image is valid.
+        Set the proper type for valid images."""
+        format_map = [
+            ("qcow", "qcow2-norm"),
+            ("iso", "iso")
+        ]
+        with magic.Magic() as m:
+            ftype = m.id_filename(self.get_path())
+            logger.debug("Downloaded file type is: %s", ftype)
+            for file_type, disk_format in format_map:
+                if file_type in ftype.lower():
+                    self.format = disk_format
+                    return True
+        return False
+
     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)
         r = requests.get(url, stream=True)
-        if r.status_code == 200:
+        if r.status_code != 200:
+            raise Exception("Invalid response status code: %s" % r.status_code)
+
         class AbortException(Exception):
             pass
         if task.is_aborted():
@@ -178,6 +197,10 @@ class Disk(object):
                     meta={'size': actsize, 'extracting': 'zip',
                           'percent': 99})
                 self.extract_iso_from_zip(disk_path)
+            if not self.check_valid_image():
+                os.unlink(disk_path)
+                raise Exception("Invalid file format. Only qcow and "
+                                "iso files are allowed.")
 
     def extract_iso_from_zip(self, disk_path):
         with ZipFile(disk_path, 'r') as z:

storagedriver.py

@@ -36,7 +36,8 @@ class download(AbortableTask):
         parent_id = kwargs.get("parent_id", None)
         disk = Disk.deserialize(disk_desc)
         disk.download(self, url, parent_id)
-        return disk.size
+        return {'size': disk.size,
+                'type': disk.format}
 
 
 @celery.task()
@@ -89,7 +90,7 @@ def move_to_trash(datastore, disk_name):
     disk_path = path.join(datastore, disk_name)
     if not path.isdir(trash_path):
         mkdir(trash_path)
-    #TODO: trash dir configurable?
+    # TODO: trash dir configurable?
     move(disk_path, trash_path)