Add import disk feature

disk.py

@@ -266,14 +266,29 @@ class Disk(object):
                 raise Exception("Invalid file format. Only qcow and "
                                 "iso files are allowed. Image from: %s" % url)
 
+    def import_disk(self, url):
+        r = requests.get(url, stream=True)
+        downloaded_file = os.path.join(self.dir, url.split('/')[-1])
+        with open(downloaded_file, 'wb') as f:
+            f.write(r.content)
+
+        cmdline = ['qemu-img',
+                   'convert',
+                   '-O', 'qcow2',
+                   downloaded_file,
+                   self.get_path()]
+        subprocess.check_output(cmdline)
+
+        os.unlink(downloaded_file)
+
+        if not self.check_valid_image():
+            os.unlink(self.get_path())
+            raise Exception("Invalid file format.")
+
+        self.size = Disk.get(self.dir, self.name).size
+
     def export(self, format, exported_name, upload_link):
-        format_dict = {
-            'vmdk': 'vmdk',
-            'qcow2': 'qcow2',
-            'vdi': 'vdi',
-            'vpc': 'vhd',
-        }
-        exported_path = self.get_path() + '.' + format_dict[format]
+        exported_path = self.get_path() + '.' + format
         cmdline = ['qemu-img',
                    'convert',
                    '-O', format,
@@ -285,14 +300,15 @@ class Disk(object):
         with open(exported_path, 'rb') as exported_disk:
             try:
                 m = MultipartEncoder(
-                    {'data': (exported_name + '.' + format_dict[format],
+                    {'data': (exported_name + '.' + format,
                               exported_disk)}
                 )
                 response = requests.post(
                     upload_link,
                     data=m,
                     headers={'Content-Type': m.content_type},
-                    params={'no_redirect': ''})
+                    params={'no_redirect': ''}
+                )
                 if response.status_code != 200:
                     raise Exception("Invalid response status code: %s" %
                                     response.status_code)

storagedriver.py

@@ -42,6 +42,13 @@ class download(AbortableTask):
 
 
 @celery.task()
+def import_disk(disk_desc, url):
+    disk = Disk.deserialize(disk_desc)
+    disk.import_disk(url)
+    return disk.size
+
+
+@celery.task()
 def export(disk_desc, format, exported_name, upload_link):
     disk = Disk.deserialize(disk_desc)
     disk.export(format, exported_name, upload_link)