Merge branch 'devel' of https://github.com/mbusb/multibootusb.git into devel

scripts/config.py

@@ -35,6 +35,7 @@ imager_usb_disk_selected = ""
 imager_lock = ""
 imager_percentage = ""
 imager_status_text = ""
+imager_return = ""
 
 install_size = ""
 

scripts/imager.py

@@ -7,32 +7,43 @@
 # under the terms of GNU General Public License, v.2 or above
 # WARNING : Any boot-able USB made using this module will destroy data stored on USB disk.
 
-import os
-import subprocess
 import collections
+import io
+import os
 import platform
 import signal
+import time
+import subprocess
+import traceback
+
 from PyQt5 import QtWidgets
 from .gui.ui_multibootusb import Ui_MainWindow
 from .gen import *
-from . import iso
 from . import config
+from . import iso
+from . import osdriver
 from . import progressbar
+from . import osdriver
+from . import udisks
+from . import usb
+
 
 if platform.system() == "Windows":
     import win32com.client
 
-
-def dd_linux():
-    import time
-    _input = "if=" + config.image_path
-    in_file_size = float(os.path.getsize(config.image_path))
-    _output = "of=" + config.usb_disk
-    os.system("umount " + config.usb_disk + "1")
-    command = ['dd', _input, _output, "bs=1M", "oflag=sync"]
-    log("Executing ==> " + " ".join(command))
-    dd_process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False)
-
+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())
+
+def _dd_iso_image(dd_progress_thread):
     pbar = progressbar.ProgressBar(
             maxval=100,
             widgets=[
@@ -43,50 +54,74 @@ def dd_linux():
             ]
     ).start()
 
-    while dd_process.poll() is None:
-        time.sleep(0.1)  # If this time delay is not given, the Popen does not execute the actual command
-        dd_process.send_signal(signal.SIGUSR1)
-        dd_process.stderr.flush()
-        while True:
-            time.sleep(0.1)
-            out_error = dd_process.stderr.readline().decode()
-            if out_error:
-                if 'bytes' in out_error:
-                    copied = int(out_error.split(' ', 1)[0])
-                    config.imager_percentage = round((float(copied) / float(in_file_size) * 100))
-                    pbar.update(config.imager_percentage)
-                    break
-
-    if dd_process.poll() is not None:
-        log("\nExecuting ==> sync")
-        os.sync()
-        log("ISO has been written to USB disk...")
-        return
+
+    def gui_update(percentage):
+        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)]
+    really_unmounted = []
+    try:
+        for c, pname in unmounted_contexts:
+                c.__enter__()
+                really_unmounted.append((c, pname))
+        error = osdriver.dd_iso_image(
+            config.image_path, config.usb_disk, gui_update)
+        if error:
+            dd_progress_thread.set_error(error)
+    finally:
+        for c, pname 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 = resource_path(os.path.join("data", "tools", "dd", "dd.exe"))
-    if os.path.exists(resource_path(os.path.join("data", "tools", "dd", "dd.exe"))):
-        log("dd exist")
-    _input = "if=" + config.image_path
+    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=\\\.\\" + config.usb_disk
-    command = [windd, _input, _output, "bs=1M", "--progress"]
-    log("Executing ==> " + " ".join(command))
+    _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():
-                log("Error writing to disk...")
+                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))
 
-        log("ISO has been written to USB disk...")
+        if config.imager_return is False:
+          log("Error writing to disk...")
+        else:
+          log("ISO has been written to USB disk...")
 
         return
 

scripts/install.py

@@ -198,11 +198,14 @@ def replace_syslinux_modules(syslinux_version, under_this_dir):
                 try:
                     with lzma.open(dst_path) as f:
                         expanded = f.read()
-                except (OSError, IOError, lzma.LZMAError):
+                except lzma.LZMAError:
+                    continue
+                except (OSError, IOError) as e:
+                    log("%s while accessing %s." % (e, dst_path))
                     continue
                 with open(dst_path, 'wb') as f:
                     f.write(expanded)
-                log("Successfully dcompressed %s." % fname)
+                log("Successfully decompressed %s." % fname)
                 continue
             try:
                 os.remove(dst_path)
@@ -230,13 +233,10 @@ def install_patch():
         config.usb_mount, "multibootusb", iso_basename(config.image_path))
     config.syslinux_version = isolinux_version(isolinux_path)
 
-    if config.distro == 'slitaz':
+    if config.distro in ['slitaz', 'ubunu']:
         replace_syslinux_modules(config.syslinux_version, distro_install_dir)
-        c32box_path = os.path.join(distro_install_dir, 'boot', 'isolinux',
-                                   'c32box.c32')
     elif config.distro == 'gentoo':
         replace_syslinux_modules(config.syslinux_version, distro_install_dir)
-
     elif config.distro == 'debian':
         iso_file_list = iso.iso_file_list(config.image_path)
         if not any(s.strip().lower().endswith("makeboot.sh")

scripts/mbusb_gui.py

@@ -31,7 +31,7 @@ from .distro import *
 from .qemu import *
 from .iso import *
 # from .imager import *
-from .imager import Imager, dd_linux, dd_win
+from .imager import Imager, dd_iso_image
 from . import persistence
 from . import config
 from . import admin
@@ -731,8 +731,15 @@ Proceed with installation?'''.lstrip() % \
 		config.process_exist = None
 
 		msgBox = QtWidgets.QMessageBox()
-		msgBox.setText("Image succesfully written to USB disk.")
-		msgBox.setInformativeText("Reboot to boot from USB or test it from <b>Boot ISO/USB</b> tab.");
+		if self.progress_thread_dd.error:
+			title = "Failed to write the iso image to the USB disk."
+			msg = self.progress_thread_dd.error
+		else:
+			title = "Image succesfully written to USB disk."
+			msg = "Reboot to boot from USB or test it from " \
+			  "<b>Boot ISO/USB</b> tab."
+		msgBox.setText(title)
+		msgBox.setInformativeText(msg);
 		msgBox.setStandardButtons(QtWidgets.QMessageBox.Ok)
 		msgBox.setIcon(QtWidgets.QMessageBox.Information)
 		msgBox.exec_()
@@ -963,15 +970,15 @@ class DD_Progress(QtCore.QThread):
 
 	def __init__(self):
 		QtCore.QThread.__init__(self)
-
-		if platform.system() == 'Linux':
-			self.thread = GenericThread(dd_linux)
-		elif platform.system() == 'Windows':
-			self.thread = GenericThread(dd_win)
+		self.error = None
+		self.thread = GenericThread(partial(dd_iso_image, self))
 
 	def __del__(self):
 		self.wait()
 
+	def set_error(self, error):
+		self.error = error
+
 	def run(self):
 		self.thread.start()
 		while self.thread.isRunning():

scripts/osdriver.py

@@ -2,9 +2,12 @@ import logging
 import logging.handlers
 import os
 import platform
+import queue
 import shutil
+import signal
 import subprocess
 import sys
+import time
 
 
 def log(message, info=True, error=False, debug=False, _print=True):
@@ -59,6 +62,8 @@ def get_physical_disk_number(usb_disk):
     :param usb_disk: USB disk (Like F:)
     :return: Disk number.
     """
+    import pythoncom
+    pythoncom.CoInitialize()
     partition, logical_disk = wmi_get_drive_info(usb_disk)
     log("Physical Device Number is %d" % partition.DiskIndex)
     return partition.DiskIndex
@@ -136,6 +141,29 @@ class Base:
         else:
             log("%s succeeded." % str(cmd))
 
+
+    def dd_iso_image(self, input_, output, gui_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)
+        dd_process = subprocess.Popen(cmd, **kw_args)
+        errors = queue.Queue()
+        while dd_process.poll() is None:
+            self.dd_iso_image_readoutput(dd_process, gui_update, in_file_size,
+                                         errors)
+        error_list = [errors.get() for i in range(errors.qsize())]
+        return self.dd_iso_image_interpret_result(
+            dd_process.returncode, error_list)
+
 class Windows(Base):
 
     def __init__(self):
@@ -144,6 +172,39 @@ class Windows(Base):
     def dd_add_args(self, cmd_vec, input, output, bs, count):
         pass
 
+    def dd_iso_image_add_args(self, cmd_vec, input_, output):
+        cmd_vec.append('--progress')
+
+    def add_dd_iso_image_popen_args(self, dd_iso_image_popen_args):
+        dd_iso_image_popen_args['universal_newlines'] = True
+
+    def dd_iso_image_readoutput(self, dd_process, gui_update, in_file_size,
+                                error_log):
+        for line in iter(dd_process.stderr.readline, ''):
+            line = line.strip()
+            if line:
+                l = line.replace(',', '')
+                if l[-1:] == 'M':
+                    bytes_copied = float(l.rstrip('M')) * 1024 * 1024
+                elif l.isdigit():
+                    bytes_copied = float(l)
+                else:
+                    if 16 < error_log.qsize():
+                        error_log.get()
+                    error_log.put(line)
+                    continue
+                gui_update(bytes_copied / in_file_size * 100.)
+                continue
+        # Now the 'dd' process should have completed or going to soon.
+
+    def dd_iso_image_interpret_result(self, returncode, error_list):
+        # dd.exe always returns 0
+        if any([ 'invalid' in s or 'error' in s for s
+                 in [l.lower() for l in error_list] ]):
+            return '\n'.join(error_list)
+        else:
+            return None
+
     def physical_disk(self, usb_disk):
         return r'\\.\physicaldrive%d' % get_physical_disk_number(usb_disk)
 
@@ -158,6 +219,37 @@ class Linux(Base):
     def dd_add_args(self, cmd_vec, input, output, bs, count):
         cmd_vec.append('conv=notrunc')
 
+    def dd_iso_image_add_args(self, cmd_vec, input_, output):
+        cmd_vec.append('oflag=sync')
+
+    def add_dd_iso_image_popen_args(self, dd_iso_image_popen_args):
+        pass
+
+    def dd_iso_image_readoutput(self, dd_process, gui_update, in_file_size,
+                                error_log):
+        # If this time delay is not given, the Popen does not execute
+        # the actual command
+        time.sleep(0.1)
+        dd_process.send_signal(signal.SIGUSR1)
+        dd_process.stderr.flush()
+        while True:
+            time.sleep(0.1)
+            out_error = dd_process.stderr.readline().decode()
+            if out_error:
+                if 'bytes' in out_error:
+                    bytes_copied = float(out_error.split(' ', 1)[0])
+                    gui_update( bytes_copied / in_file_size * 100. )
+                    break
+                if 16 < error_log.qsize():
+                    error_log.get()
+                error_log.put(out_error)
+            else:
+                # stderr is closed
+                break
+
+    def dd_iso_image_interpret_result(self, returncode, error_list):
+        return None if returncode==0 else '\n'.join(error_list)
+
     def physical_disk(self, usb_disk):
         return usb_disk.rstrip('0123456789')
 
@@ -177,6 +269,7 @@ for func_name in [
         'run_dd',
         'physical_disk',
         'mbusb_log_file',
+        'dd_iso_image',
         ]:
     globals()[func_name] = getattr(osdriver, func_name)
 

scripts/syslinux.py

@@ -198,6 +198,12 @@ def syslinux_default(usb_disk):
                 '''
                 if config.usb_gpt is False:
                     log('\nExecuting ==> ' + mbr_install_cmd)
+                    #
+                    # Updating mbr using dd results in catastrophy.
+                    # Windows will lose track of the filesystem and
+                    # the target volume will go away.
+                    # Never enable this code without proper work around!
+                    #
                     if subprocess.call(mbr_install_cmd, shell=True) == 0:
                         log("\nmbr install is success...\n")
                         return True

scripts/udisks.py

@@ -15,18 +15,25 @@ import os
 import re
 
 
-def node_mountpoint(node):
+def de_mangle_mountpoint(raw):
+    return raw.replace('\\040', ' ').replace('\\011', '\t') \
+      .replace('\\012', '\n').replace('\\0134', '\\')
 
-    def de_mangle(raw):
-        return raw.replace('\\040', ' ').replace('\\011', '\t').replace('\\012',
-                '\n').replace('\\0134', '\\')
+def node_mountpoint(node):
 
     for line in open('/proc/mounts').readlines():
         line = line.split()
         if line[0] == node:
-            return de_mangle(line[1])
+            return de_mangle_mountpoint(line[1])
     return None
 
+def find_partitions_on(disk):
+    assert not disk[-1:].isdigit()
+    with open('/proc/mounts') as f:
+        relevant_lines = [l.split(' ') for l in f.readlines()
+                          if l.startswith(disk)]
+    return [ [v[0], de_mangle_mountpoint(v[1])] + v[2:] for v
+             in relevant_lines ]
 
 class NoUDisks1(Exception):
     pass

scripts/usb.py

@@ -476,7 +476,7 @@ class UnmountedContext:
         gen.log("Unmounted %s" % self.usb_disk)
         return self
 
-    def __exit__(self, type, value, traceback_):
+    def __exit__(self, type_, value, traceback_):
         if not self.is_relevant:
             return
         os.sync()     # This should not be strictly necessary