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

multibootusb

@@ -101,7 +101,7 @@ Example for writing ISO image to target USB disk (will destroy data on USB disk)
     Windows:
         python3 multibootusb -c -i -r ../../favourite.iso -t G:
 ''')
-    exit(2)
+    sys.exit(2)
 
 
 def start_gui():
@@ -121,9 +121,9 @@ if __name__ == '__main__':
             admin.runAsAdmin()
             sys.exit(0)
     try:
-        opts, args = getopt.getopt(sys.argv[1:], 'i:t:yvhcudr',
+        opts, args = getopt.getopt(sys.argv[1:], 'i:t:yvhcudrs',
                                    ['iso=', 'target=', 'yes', 'version', 'help', 'command', 'uninstall', 'debug',
-                                    'raw'])
+                                    'raw', 'syslinux'])
     except getopt.GetoptError:
         usage()
         sys.exit(2)
@@ -151,6 +151,8 @@ if __name__ == '__main__':
             config.yes = True
         elif opt in ('-r', '--raw'):
             config.cli_dd = True
+        elif opt in ('-s', '--syslinux'):
+            config.cli_syslinux = True
         else:
             gui = True
             #start_gui()
@@ -174,6 +176,8 @@ if gui is False:
     elif config.image_path is '' and config.usb_disk is '':
         log('\nNo option provided. See the usage below.')
         usage()
+    elif config.cli_syslinux is True and config.usb_disk is not '':
+        cli_install_syslinux()
     elif config.image_path is '' or config.usb_disk is '':
         log('\nOptions \'-i\' and \'-t\' must be supplied together. See the usage below.')
         usage()

scripts/config.py

@@ -27,6 +27,7 @@ iso_bin_dir = ''
 process_exist = None
 yes = False
 cli_dd = False
+cli_syslinux = False
 
 imager_iso_link = ""
 imager_usb_disk_selected = ""

scripts/gen.py

@@ -229,6 +229,23 @@ def copy_mbusb_dir_usb(usb_disk):
     else:
         log('EFI directory already exist. Not copying.')
 
+    # For backward compatibility
+    if not os.path.exists(os.path.join(usb_mount_path, 'EFI', 'BOOT', 'bootx64-gpt.efi')):
+        shutil.copy(resource_path(os.path.join('data', 'EFI', 'BOOT', 'bootx64-gpt.efi')),
+                    os.path.join(usb_mount_path, 'EFI', 'BOOT', 'bootx64-gpt.efi'))
+
+    if not os.path.exists(os.path.join(usb_mount_path, 'EFI', 'BOOT', 'bootx64-msdos.efi')):
+        shutil.copy(resource_path(os.path.join('data', 'EFI', 'BOOT', 'bootx64-msdos.efi')),
+                    os.path.join(usb_mount_path, 'EFI', 'BOOT', 'bootx64-msdos.efi'))
+
+    if not os.path.exists(os.path.join(usb_mount_path, 'multibootusb', 'grub', 'core-gpt.img')):
+        shutil.copy(resource_path(os.path.join('data', 'multibootusb', 'grub', 'core-gpt.img')),
+                    os.path.join(usb_mount_path, 'multibootusb', 'grub', 'core-gpt.img'))
+
+    if not os.path.exists(os.path.join(usb_mount_path, 'multibootusb', 'grub', 'core-msdos.img')):
+        shutil.copy(resource_path(os.path.join('data', 'multibootusb', 'grub', 'core-msdos.img')),
+                    os.path.join(usb_mount_path, 'multibootusb', 'grub', 'core-msdos.img'))
+
     return result
 
 

scripts/mbusb_cli.py

@@ -16,6 +16,7 @@ from .distro import *
 from .syslinux import *
 from .install import *
 from . import imager
+from . import syslinux
 
 
 def read_input_uninstall():
@@ -99,6 +100,7 @@ def iso_install(iso_image):
                     install_progress()
                     syslinux_distro_dir(config.usb_disk, iso_image, _distro)
                     syslinux_default(config.usb_disk)
+                    replace_grub_binary()
                     update_distro_cfg_files(iso_image, config.usb_disk, _distro)
                     log('Finished installing ' + iso.iso_basename(iso_image))
                 else:
@@ -110,6 +112,7 @@ def iso_install(iso_image):
                 install_progress()
                 syslinux_distro_dir(config.usb_disk, iso_image, _distro)
                 syslinux_default(config.usb_disk)
+                replace_grub_binary()
                 update_distro_cfg_files(iso_image, config.usb_disk, _distro)
                 log('Finished installing ' + iso.iso_basename(iso_image))
         else:
@@ -167,3 +170,40 @@ def cli_dd():
         else:
             log('\nAuto install is not recommended in direct writing method. Please choose without \'-y\' option.\n')
             sys.exit(2)
+
+
+def cli_install_syslinux():
+    """
+    Install syslinux on a target USB disk. It will installed on 'multibootusb' directory
+    :return: 
+    """
+    if platform.system() == 'Linux':
+        if config.usb_disk[-1].isdigit() is not True:
+            log('Selected USB disk is not a partition. Please enter the partition eg. \'/dev/sdb1\'')
+            sys.exit(2)
+        elif is_root() is False:
+            log("You need to have root privileges to run this script.\nPlease try again using admin privilege (sudo).")
+            sys.exit(2)
+
+    if config.yes is not True:
+        log('\nInitiating process for installing syslinux on ' + config.usb_disk)
+        log('Selected target device is  : ' + quote(config.usb_disk))
+        log('Syslinux install directory :  \'multibootusb\'\n')
+        log('Please confirm the option.')
+        log('Y/y/Yes/yes/YES or N/n/No/no/NO')
+        if read_input_yes() is True:
+            if syslinux.syslinux_default(config.usb_disk) is True:
+                log('Syslinux successfully installed on ' + config.usb_disk)
+            else:
+                log('Failed to install syslinux on ' + config.usb_disk)
+        else:
+            log('Operation cancelled by user. Exiting...')
+            sys.exit(2)
+    else:
+        log('\nSkipping user input and installing syslinux on ' + config.usb_disk)
+        if syslinux.syslinux_default(config.usb_disk) is True:
+            log('Syslinux successfully installed on ' + config.usb_disk)
+        else:
+            log('Failed to install syslinux on ' + config.usb_disk)
+        sys.exit(2)
+

scripts/mbusb_gui.py

@@ -309,6 +309,7 @@ Are you SURE you want to enable it?",
         self.ui.statusbar.showMessage(str("Status: Installing Syslinux..."))
         syslinux_distro_dir(config.usb_disk, config.image_path, config.distro)
         syslinux_default(config.usb_disk)
+        replace_grub_binary()
         update_distro_cfg_files(config.image_path, config.usb_disk, config.distro, config.persistence)
         self.update_list_box(config.usb_disk)
         if sys.platform.startswith("linux"):

scripts/syslinux.py

@@ -21,21 +21,68 @@ syslinux_fs = ["vfat", "ntfs", "FAT32", "NTFS"]
 mbr_bin = resource_path(os.path.join("data", "tools", "mbr.bin"))
 
 
+def gpt_part_table(usb_disk):
+    """
+    Check if selected USB contain GPT or MBR partition table
+    :return: True if GPT else False
+    """
+    if platform.system() == "Linux":
+        _cmd_out = subprocess.check_output("parted  " + usb_disk[:-1] + " print", shell=True)
+        if b'msdos' in _cmd_out:
+            return False
+        elif b'gpt' in _cmd_out:
+            return True
+    elif platform.system() == 'Windows':
+        _cmd_out = ''
+
+
+def get_mbr_bin_path(usb_disk):
+    """
+    Check if partition table type is mbr or gpr using parted command under Linux
+    :param usb_disk: path to whole USB disk '/dev/sdb'
+    :return: Path to mbr.bin for use
+    """
+    if platform.system() == "Linux":
+        if gpt_part_table(usb_disk) is False:
+            log('Using mbr.bin msdos mbr install.')
+            return resource_path(os.path.join("data", "tools", "mbr.bin"))
+        if gpt_part_table(usb_disk) is True:
+            log('Using gptmbr.bin for mbr install.')
+            return resource_path(os.path.join("data", "tools", "gptmbr.bin"))
+    elif platform.system() == 'Windows':
+        _cmd_out = ''
+
+    return False
+
+
 def set_boot_flag(usb_disk):
     if platform.system() == "Linux":
         log("\nChecking boot flag on " + usb_disk[:-1], '\n')
         cmd_out = subprocess.check_output("parted -m -s " + usb_disk[:-1] + " print", shell=True)
-        if b'boot' in cmd_out:
-            log("\nDisk " + usb_disk[:-1] + " already has boot flag.\n")
-            return True
-        else:
-            log("\nExecuting ==>  parted " + usb_disk[:-1] + " set 1 boot on", '\n')
-            if subprocess.call("parted " + usb_disk[:-1] + " set 1 boot on", shell=True) == 0:
-                log("\nBoot flag set to bootable " + usb_disk[:-1], '\n')
+        if gpt_part_table(usb_disk) is False:
+            if b'boot' in cmd_out:
+                log("\nDisk " + usb_disk[:-1] + " already has boot flag.\n")
                 return True
             else:
-                log("\nUnable to set boot flag on  " + usb_disk[:-1], '\n')
-                return False
+                log("\nExecuting ==>  parted " + usb_disk[:-1] + " set 1 boot on", '\n')
+                if subprocess.call("parted " + usb_disk[:-1] + " set 1 boot on", shell=True) == 0:
+                    log("\nBoot flag set to bootable " + usb_disk[:-1], '\n')
+                    return True
+                else:
+                    log("\nUnable to set boot flag on  " + usb_disk[:-1], '\n')
+                    return False
+        elif gpt_part_table(usb_disk) is True:
+            if b'legacy_boot' in cmd_out:
+                log("\nGPT Disk " + usb_disk[:-1] + " already has legacy_boot flag.\n")
+                return True
+            else:
+                log("\nExecuting ==>  parted " + usb_disk[:-1] + " set 1 legacy_boot on", '\n')
+                if subprocess.call("parted " + usb_disk[:-1] + " set 1 legacy_boot on", shell=True) == 0:
+                    log("\nBoot flag set to legacy_boot " + usb_disk[:-1], '\n')
+                    return True
+                else:
+                    log("\nUnable to set legacy_boot flag on  " + usb_disk[:-1], '\n')
+                    return False
 
 
 def syslinux_default(usb_disk):
@@ -48,6 +95,7 @@ def syslinux_default(usb_disk):
     usb_details = usb.details(usb_disk)
     usb_fs = usb_details['file_system']
     usb_mount = usb_details['mount_point']
+    mbr_bin = get_mbr_bin_path(usb_disk)
     mbr_install_cmd = 'dd bs=440 count=1 conv=notrunc if=' + mbr_bin + ' of=' + usb_disk[:-1]
     # log(usb_fs)
     if usb_fs in extlinux_fs:
@@ -78,6 +126,7 @@ def syslinux_default(usb_disk):
             if subprocess.call(syslinux_cmd, shell=True) == 0:
                 log("\nDefault syslinux install is success...\n")
                 config.status_text = 'Default syslinux successfully installed...'
+                log('\nExecuting ==> ' + mbr_install_cmd)
                 if subprocess.call(mbr_install_cmd, shell=True) == 0:
                     config.status_text = 'mbr install is success...'
                     log("\nmbr install is success...\n")
@@ -203,6 +252,50 @@ def syslinux_distro_dir(usb_disk, iso_link, distro):
                     else:
                         log("\nFailed to install syslinux on distro directory...\n")
 
+
+def replace_grub_binary():
+    """
+    This function checks if correct binary is installed on grub and EFI directory.
+    If mismatch is found between partition table and binary, replace it correct one.
+    Default binaries will work for msdos partition table and therefore need not be replaced.
+    :return: 
+    """
+    grub_efi_bin_path = os.path.join(config.usb_mount, 'EFI', 'BOOT', 'bootx64.efi')
+    grub_efi_bin_gpt_path = os.path.join(config.usb_mount, 'EFI', 'BOOT', 'bootx64-gpt.efi')
+    grub_efi_bin_msdos_path = os.path.join(config.usb_mount, 'EFI', 'BOOT', 'bootx64-msdos.efi')
+    core_img_bin_path = os.path.join(config.usb_mount, 'multibootusb', 'grub', 'core.img')
+    core_img_bin_msdos_path = os.path.join(config.usb_mount, 'multibootusb', 'grub', 'core-msdos.img')
+    core_img_bin_gpt_path = os.path.join(config.usb_mount, 'multibootusb', 'grub', 'core-gpt.img')
+
+    if platform.system() == 'Linux':
+        if gpt_part_table(config.usb_disk) is True:
+            log('Replacing efi binary with gpt compatible one...')
+            try:
+                shutil.copy(grub_efi_bin_gpt_path, grub_efi_bin_path)
+            except Exception as e:
+                log(e)
+                log('Failed to replace efi binary...')
+            log('Replacing core.img binary with gpt compatible one...')
+            try:
+                shutil.copy(core_img_bin_gpt_path, core_img_bin_path)
+            except Exception as e:
+                log(e)
+                log('Failed to replace efi binary...')
+        else:
+            log('Replacing efi binary with msdos compatible one...')
+            try:
+                shutil.copy(core_img_bin_gpt_path, core_img_bin_path)
+            except Exception as e:
+                log(e)
+                log('Failed to replace efi binary...')
+            log('Replacing efi binary with msdos compatible one...')
+            try:
+                shutil.copy(core_img_bin_msdos_path, core_img_bin_path)
+            except Exception as e:
+                log(e)
+                log('Failed to replace core.img binary...')
+
+
 if __name__ == '__main__':
     if os.geteuid() != 0:
         log('Please running this script with sudo/root/admin privilage.')