Use pyparted to check for bootflag [Closes: issue1248]

[grml2usb.git] / grml2usb

diff --git a/grml2usb b/grml2usb
index eafee4b..2f58d01 100755 (executable)
--- a/grml2usb
+++ b/grml2usb
@@ -40,12 +40,13 @@ GRML_FLAVOURS = set()  # which flavours are being installed?
 GRML_DEFAULT = None
 UUID = None
 SYSLINUX_LIBS = "/usr/lib/syslinux/"
-GPT_HEADER = "\x55\xaa\x45\x46\x49\x20\x50\x41\x52\x54" # original GPT header
+GPT_HEADER = "\x55\xaa\x45\x46\x49\x20\x50\x41\x52\x54"  # original GPT header
 
 RE_PARTITION = re.compile(r'([a-z/]*?)(\d+)$')
 RE_P_PARTITION = re.compile(r'(.*?\d+)p(\d+)$')
 RE_LOOP_DEVICE = re.compile(r'/dev/loop\d+$')
 
+
 def syslinux_warning(option, opt, value, opt_parser):
     """A helper function for printing a warning about deprecated option
     """
@@ -132,6 +133,7 @@ class CriticalException(Exception):
     @Exception: message"""
     pass
 
+
 class VerifyException(Exception):
     """Throw critical exception if there is an fatal error when verifying something.
 
@@ -287,7 +289,7 @@ def search_file(filename, search_path='/bin' + os.pathsep + '/usr/bin', lst_retu
 
         @cwd: current working directory
         """
-        return  os.path.exists(os.path.join(cwd, filename))
+        return os.path.exists(os.path.join(cwd, filename))
 
     for path in paths:
         current_dir = path
@@ -315,9 +317,29 @@ def check_uid_root():
         raise CriticalException("please run this script with uid 0 (root).")
 
 
+# for usage inside check_boot_flag
+def get_partition_for_path(path):
+    import parted
+
+    boot_dev, x = get_device_from_partition(path)
+
+    d = parted.getDevice(boot_dev)
+    disk = parted.Disk(d)
+    return disk.getPartitionByPath(path)
+
+
 def check_boot_flag(device):
     boot_dev, x = get_device_from_partition(device)
 
+    try:
+        import parted
+        part = get_partition_for_path(device)
+        if part.getFlag(parted.PARTITION_BOOT):
+            logging.debug("bootflag is enabled on %s" % device)
+            return
+    except ImportError, e:
+        pass
+
     with open(boot_dev, 'r') as image:
         data = image.read(520)
         bootcode = data[440:]
@@ -330,7 +352,7 @@ def check_boot_flag(device):
         else:
             logging.debug("bootflag is NOT enabled")
             raise VerifyException("Device %s does not have the bootflag set. "
-                "Please enable it to be able to boot." % boot_dev)
+                "Please enable it to be able to boot." % device)
 
 
 def mkfs_fat16(device):
@@ -801,8 +823,7 @@ def copy_system_files(grml_flavour, iso_mount, target):
     squashfs = search_file(grml_flavour + '.squashfs', iso_mount)
     if squashfs is None:
         logging.error("error locating squashfs file")
-        raise CriticalException("squashfs file not found"
-            ", please check that your iso is not corrupt")
+        raise CriticalException("squashfs file not found, please check that your iso is not corrupt")
     else:
         squashfs_target = target + '/live/' + grml_flavour + '/'
         execute(mkdir, squashfs_target)