From 7f3fe377c34231ab9b6565b5a3e9cad1aa865b97 Mon Sep 17 00:00:00 2001 From: Shinji Suzuki Date: Sat, 19 May 2018 04:40:18 +0900 Subject: [PATCH] Unify imager code again. Disable assignment to 'config.imager_return' as it is not used elsewhere yet. --- scripts/gen.py | 17 +-------- scripts/imager.py | 93 ++++++++++++--------------------------------- scripts/osdriver.py | 47 +++++++++++++++++++++-- scripts/udisks.py | 2 +- 4 files changed, 70 insertions(+), 89 deletions(-) diff --git a/scripts/gen.py b/scripts/gen.py index 46e6e57..29f6454 100644 --- a/scripts/gen.py +++ b/scripts/gen.py @@ -13,13 +13,12 @@ import platform import shutil import string import zipfile -import tempfile import re import ctypes from . import config from .osdriver import get_physical_disk_number, wmi_get_drive_info, \ - log, resource_path + log, resource_path, multibootusb_host_dir def scripts_dir_path(): return os.path.dirname(os.path.realpath(__file__)) @@ -71,20 +70,6 @@ def sys_64bits(): return sys.maxsize > 2**32 -def multibootusb_host_dir(): - """ - Cross platform way to detect multibootusb directory on host system. - :return: Path to multibootusb directory of host system. - """ - if platform.system() == "Linux": - home_dir = os.path.expanduser('~') - mbusb_dir = os.path.join(home_dir, ".multibootusb") - elif platform.system() == "Windows": - mbusb_dir = os.path.join(tempfile.gettempdir(), "multibootusb") - - return mbusb_dir - - def iso_cfg_ext_dir(): """ Function to return the path to ISO configuration file extraction directory. diff --git a/scripts/imager.py b/scripts/imager.py index 763360c..05f1164 100644 --- a/scripts/imager.py +++ b/scripts/imager.py @@ -24,7 +24,6 @@ from . import iso from . import osdriver from . import progressbar from . import osdriver -from . import udisks from . import usb @@ -32,16 +31,14 @@ if platform.system() == "Windows": import win32com.client def dd_iso_image(dd_progress_thread): - if platform.system() == "Windows": - dd_win() - else: - try: - _dd_iso_image(dd_progress_thread) - except: - o = io.StringIO() - traceback.print_exc(None, o) - log(o.getvalue()) - dd_progress_thread.set_error(o.getvalue()) + try: + _dd_iso_image(dd_progress_thread) + except: + # config.imager_return = False + o = io.StringIO() + traceback.print_exc(None, o) + log(o.getvalue()) + dd_progress_thread.set_error(o.getvalue()) def _dd_iso_image(dd_progress_thread): pbar = progressbar.ProgressBar( @@ -59,73 +56,31 @@ def _dd_iso_image(dd_progress_thread): config.imager_percentage = percentage pbar.update(percentage) - unmounted_contexts = [ - (usb.UnmountedContext(p[0], config.update_usb_mount), p[0]) for p - in udisks.find_partitions_on(config.usb_disk)] + def status_update(text): + config.status_text = text + + mounted_partitions = osdriver.find_mounted_partitions_on(config.usb_disk) really_unmounted = [] try: - for c, pname in unmounted_contexts: - c.__enter__() - really_unmounted.append((c, pname)) + for x in mounted_partitions: + partition_dev, mount_point = x[:2] + c = usb.UnmountedContext(partition_dev, config.update_usb_mount) + c.__enter__() + really_unmounted.append((c, partition_dev)) error = osdriver.dd_iso_image( - config.image_path, config.usb_disk, gui_update) + config.image_path, config.usb_disk, gui_update, status_update) if error: dd_progress_thread.set_error(error) + log('Error writing iso image...') + # config.imager_return = False + else: + log('ISO has been written to USB disk...') + # config.imager_return = True finally: - for c, pname in really_unmounted: + for c, partition_dev in really_unmounted: c.__exit__(None, None, None) -def dd_win_clean_usb(usb_disk_no): - """ - - """ - host_dir = multibootusb_host_dir() - temp_file = os.path.join(host_dir, "preference", "disk_part.txt") - diskpart_text_feed = 'select disk ' + str(usb_disk_no) + '\nclean\nexit' - write_to_file(temp_file, diskpart_text_feed) - config.status_text = 'Cleaning the disk...' - if subprocess.call('diskpart.exe -s ' + temp_file ) == 0: - return True - else: - return False - - -def dd_win(): - - windd = quote(resource_path(os.path.join("data", "tools", "dd", "dd.exe"))) - usb_disk_no = osdriver.get_physical_disk_number(config.usb_disk) - _input = "if=" + config.image_path.strip() - in_file_size = float(os.path.getsize(config.image_path) / 1024 / 1024) - _output = "of=\\\.\\PhysicalDrive" + str(usb_disk_no) - if dd_win_clean_usb(usb_disk_no) is False: - return - # _output = "od=" + config.usb_disk # 'od=' option should also work. - # command = [windd, _input, _output, "bs=1M", "--progress"] - command = windd + ' ' + _input + ' ' + _output + " bs=1M" + " --progress" - #log("Executing ==> " + " ".join(command)) - log("Executing ==> " + command) - dd_process = subprocess.Popen(command, universal_newlines=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE, - shell=False) - time.sleep(0.1) - while dd_process.poll() is None: - for line in iter(dd_process.stderr.readline, ''): - line = line.strip() - if 'error' in line.lower() or 'invalid' in line.lower(): - config.imager_return = False - break - if line and line[-1] == 'M': - copied = float(line.strip('M').replace(',', '')) - config.imager_percentage = round((copied / float(in_file_size) * 100)) - - if config.imager_return is False: - log("Error writing to disk...") - else: - log("ISO has been written to USB disk...") - - return - - class Imager(QtWidgets.QMainWindow, Ui_MainWindow): """ Raw write to USB disk using dd. diff --git a/scripts/osdriver.py b/scripts/osdriver.py index b562466..06de90c 100644 --- a/scripts/osdriver.py +++ b/scripts/osdriver.py @@ -7,8 +7,10 @@ import shutil import signal import subprocess import sys +import tempfile import time +from . import udisks def log(message, info=True, error=False, debug=False, _print=True): """ @@ -130,6 +132,26 @@ def wmi_get_drive_info_ex(usb_disk): return r +def dd_win_clean_usb(usb_disk_no, status_update): + """ + + """ + host_dir = multibootusb_host_dir() + temp_file = os.path.join(host_dir, "preference", "disk_part.txt") + diskpart_text_feed = 'select disk ' + str(usb_disk_no) + '\nclean\nexit' + with open(temp_file, 'wb') as fp: + fp.write(bytes(diskpart_text_feed, 'utf-8')) + status_update('Cleaning the disk...') + if subprocess.call('diskpart.exe -s ' + temp_file ) == 0: + for i in range(40): + # Make sure the drive reappears if it has ever gone away. + time.sleep(0.25) + with open('\\\\.\\PhysicalDrive%d' % usb_disk_no, 'rb'): + return + raise RuntimeError('PhysicalDrive%d is now gone!' % usb_disk_no) + raise RuntimeError("Execution dd(.exe) has failed.") + + class Base: def run_dd(self, input, output, bs, count): @@ -142,19 +164,19 @@ class Base: log("%s succeeded." % str(cmd)) - def dd_iso_image(self, input_, output, gui_update): + def dd_iso_image(self, input_, output, gui_update, status_update): in_file_size = os.path.getsize(input_) - cmd = [self.dd_exe, 'if=' + input_, 'of=' + self.physical_disk(output), 'bs=1M'] self.dd_iso_image_add_args(cmd, input_, output) - log('Executing => ' + str(cmd)) kw_args = { 'stdout' : subprocess.PIPE, 'stderr' : subprocess.PIPE, 'shell' : False, } self.add_dd_iso_image_popen_args(kw_args) + self.dd_iso_image_prepare(input, output, status_update) + log('Executing => ' + str(cmd)) dd_process = subprocess.Popen(cmd, **kw_args) errors = queue.Queue() while dd_process.poll() is None: @@ -172,6 +194,9 @@ class Windows(Base): def dd_add_args(self, cmd_vec, input, output, bs, count): pass + def dd_iso_image_prepare(self, input, output, status_update): + return dd_win_clean_usb(get_physical_disk_number(output), status_update) + def dd_iso_image_add_args(self, cmd_vec, input_, output): cmd_vec.append('--progress') @@ -211,11 +236,20 @@ class Windows(Base): def mbusb_log_file(self): return os.path.join(os.getcwd(), 'multibootusb.log') + def find_mounted_partitions_on(self, usb_disk): + return [] # No-op until UnmountedContext() get implemented for Windows + + def multibootusb_host_dir(self): + return os.path.join(tempfile.gettempdir(), "multibootusb") + class Linux(Base): def __init__(self): self.dd_exe = 'dd' + def dd_iso_image_prepare(self, input, output, status_update): + pass + def dd_add_args(self, cmd_vec, input, output, bs, count): cmd_vec.append('conv=notrunc') @@ -256,6 +290,11 @@ class Linux(Base): def mbusb_log_file(self): return '/var/log/multibootusb.log' + def find_mounted_partitions_on(self, usb_disk): + return udisks.find_mounted_partitions_on(usb_disk) + + def multibootusb_host_dir(self): + return os.path.join(os.path.expanduser('~'), ".multibootusb") driverClass = { 'Windows' : Windows, @@ -270,6 +309,8 @@ for func_name in [ 'physical_disk', 'mbusb_log_file', 'dd_iso_image', + 'find_mounted_partitions_on', + 'multibootusb_host_dir', ]: globals()[func_name] = getattr(osdriver, func_name) diff --git a/scripts/udisks.py b/scripts/udisks.py index be25ac1..8460d88 100644 --- a/scripts/udisks.py +++ b/scripts/udisks.py @@ -27,7 +27,7 @@ def node_mountpoint(node): return de_mangle_mountpoint(line[1]) return None -def find_partitions_on(disk): +def find_mounted_partitions_on(disk): assert not disk[-1:].isdigit() with open('/proc/mounts') as f: relevant_lines = [l.split(' ') for l in f.readlines()