diff --git a/disk.py b/disk.py index c3080c8..bd45b70 100644 --- a/disk.py +++ b/disk.py @@ -266,26 +266,22 @@ class Disk(object): raise Exception("Invalid file format. Only qcow and " "iso files are allowed. Image from: %s" % url) - def import_disk(self, task, url, parent_id): - r = requests.get(url, stream=True) - clen = int(r.headers.get('content-length')) - downloaded_file = os.path.join(self.dir, url.split('/')[-1]) - percent = 0 + def import_disk(self, task, url, port=None): + downloaded_file = os.path.join(self.dir, re.split('[:/]', url)[-1]) + cmdline = ['scp', '-B'] + if port is not None: + cmdline.extend(['-P', str(port)]) + cmdline.extend([downloaded_file, url]) + + proc = subprocess.Popen(cmdline) try: - with open(downloaded_file, 'wb') as f: - for chunk in r.iter_content(chunk_size=256 * 1024): - f.write(chunk) - current_size = f.tell() - new_percent = current_size * 100 / clen - if new_percent > percent: - if task.is_aborted(): - raise AbortException() - percent = new_percent - task.update_state( - task_id=parent_id, - state=task.AsyncResult(parent_id).state, - meta={'size': current_size, 'percent': percent} - ) + while proc.poll() is None: + if task.is_aborted(): + raise AbortException() + sleep(1) + + if task.is_aborted(): + raise AbortException() with magic.Magic() as m: ftype = m.id_filename(downloaded_file) diff --git a/storagedriver.py b/storagedriver.py index f06e8f2..d1913e3 100644 --- a/storagedriver.py +++ b/storagedriver.py @@ -48,9 +48,12 @@ class import_disk(AbortableTask): def run(self, **kwargs): disk_desc = kwargs["disk_desc"] url = kwargs["url"] - parent_id = kwargs["task"] + if "port" in kwargs: + port = kwargs["port"] + else: + port = None disk = Disk.deserialize(disk_desc) - disk.import_disk(self, url, parent_id) + disk.import_disk(self, url, port) return { "size": disk.size, "checksum": disk.checksum