[v3] wic: add wic tests and support setting GPT diskid

diff --git a/meta/lib/oeqa/selftest/cases/wic.py b/meta/lib/oeqa/selftest/cases/wic.py
index bb4ac23ebf..d7a9b14658 100644
--- a/meta/lib/oeqa/selftest/cases/wic.py
+++ b/meta/lib/oeqa/selftest/cases/wic.py
@@ -1905,6 +1905,42 @@  INITRAMFS_IMAGE = "core-image-initramfs-boot"
             self.assertIn("Source parameter 'fill' only works with the '--fixed-size' option, exiting.", result.output)
             self.assertNotEqual(0, result.status)
 
+    def test_diskid_on_msdos_partition(self):
+        """Test diksid on msdos partions"""
+        img = 'core-image-minimal'
+        diskid = "0xdeadbbef"
+        with NamedTemporaryFile("w", suffix=".wks") as wks:
+            wks.writelines(['bootloader --ptable msdos --diskid %s\n' % diskid,
+                            'part /boot --size=100M --active --fstype=ext4 --label boot\n'
+                            'part /     --source rootfs      --fstype=ext4 --label root\n'])
+            wks.flush()
+            cmd = "wic create %s -e %s -o %s" % (wks.name, img, self.resultdir)
+            runCmd(cmd)
+            wksname = os.path.splitext(os.path.basename(wks.name))[0]
+            out = glob(os.path.join(self.resultdir, "%s-*direct" % wksname))
+            self.assertEqual(1, len(out))
+            sysroot = get_bb_var('RECIPE_SYSROOT_NATIVE', 'wic-tools')
+            result = runCmd("%s/usr/sbin/sfdisk -l %s | grep 'Disk identifier:'" % (sysroot, out[0]))
+            self.assertEqual("Disk identifier: %s" % diskid.lower(), result.output)
+
+    def test_diskid_on_gpt_partition(self):
+        """Test diksid on gpt partions"""
+        img = 'core-image-minimal'
+        diskid = "deadbeef-cafe-babe-f00d-cec2ea4eface"
+        with NamedTemporaryFile("w", suffix=".wks") as wks:
+            wks.writelines(['bootloader --ptable gpt --diskid %s\n' % diskid,
+                            'part /boot --size=100M --active --fstype=ext4 --label boot\n'
+                            'part /     --source rootfs      --fstype=ext4 --label root\n'])
+            wks.flush()
+            cmd = "wic create %s -e %s -o %s" % (wks.name, img, self.resultdir)
+            runCmd(cmd)
+            wksname = os.path.splitext(os.path.basename(wks.name))[0]
+            out = glob(os.path.join(self.resultdir, "%s-*direct" % wksname))
+            self.assertEqual(1, len(out))
+            sysroot = get_bb_var('RECIPE_SYSROOT_NATIVE', 'wic-tools')
+            result = runCmd("%s/usr/sbin/sfdisk -l %s | grep 'Disk identifier:'" % (sysroot, out[0]))
+            self.assertEqual("Disk identifier: %s" % diskid.upper(), result.output)
+
 class ModifyTests(WicTestCase):
     def test_wic_ls(self):
         """Test listing image content using 'wic ls'"""
diff --git a/scripts/lib/wic/ksparser.py b/scripts/lib/wic/ksparser.py
index 48b5b09ddd..4ccd70dc55 100644
--- a/scripts/lib/wic/ksparser.py
+++ b/scripts/lib/wic/ksparser.py
@@ -16,6 +16,7 @@  import os
 import shlex
 import logging
 import re
+import uuid
 
 from argparse import ArgumentParser, ArgumentError, ArgumentTypeError
 
@@ -196,7 +197,7 @@  class KickStart():
         bootloader.add_argument('--configfile')
         bootloader.add_argument('--ptable', choices=('msdos', 'gpt', 'gpt-hybrid'),
                                 default='msdos')
-        bootloader.add_argument('--diskid', type=lambda x: int(x, 0))
+        bootloader.add_argument('--diskid')
         bootloader.add_argument('--timeout', type=int)
         bootloader.add_argument('--source')
 
@@ -297,6 +298,24 @@  class KickStart():
                             if append_var:
                                 self.bootloader.append = ' '.join(filter(None, \
                                                          (self.bootloader.append, append_var)))
+                            if parsed.diskid:
+                                if parsed.ptable == "msdos":
+                                    try:
+                                        self.bootloader.diskid = int(parsed.diskid, 0)
+                                    except ValueError:
+                                        err = "with --ptbale msdos only 32bit integers " \
+                                              "are allowed for --diskid. %s could not " \
+                                              "be parsed" % self.ptable
+                                        raise KickStartError(err)
+                                else:
+                                    try:
+                                        self.bootloader.diskid = uuid.UUID(parsed.diskid)
+                                    except ValueError:
+                                        err = "with --ptable %s only valid uuids are " \
+                                              "allowed for --diskid. %s could not be " \
+                                              "parsed" % (parsed.ptable, parsed.diskid)
+                                        raise KickStartError(err)
+
                         else:
                             err = "%s:%d: more than one bootloader specified" \
                                       % (confpath, lineno)
diff --git a/scripts/lib/wic/plugins/imager/direct.py b/scripts/lib/wic/plugins/imager/direct.py
index f40f033a3d..ad922cfbf1 100644
--- a/scripts/lib/wic/plugins/imager/direct.py
+++ b/scripts/lib/wic/plugins/imager/direct.py
@@ -315,7 +315,14 @@  class PartitionedImage():
                           # all partitions (in bytes)
         self.ptable_format = ptable_format  # Partition table format
         # Disk system identifier
-        if disk_id:
+        if disk_id and ptable_format in ('gpt', 'gpt-hybrid'):
+            self.disk_guid = disk_id
+        elif os.getenv('SOURCE_DATE_EPOCH'):
+            self.disk_guid = uuid.UUID(int=int(os.getenv('SOURCE_DATE_EPOCH')))
+        else:
+            self.disk_guid = uuid.uuid4()
+
+        if disk_id and ptable_format == 'msdos':
             self.identifier = disk_id
         elif os.getenv('SOURCE_DATE_EPOCH'):
             self.identifier = random.Random(int(os.getenv('SOURCE_DATE_EPOCH'))).randint(1, 0xffffffff)
@@ -545,11 +552,6 @@  class PartitionedImage():
 
     def _write_disk_guid(self):
         if self.ptable_format in ('gpt', 'gpt-hybrid'):
-            if os.getenv('SOURCE_DATE_EPOCH'):
-                self.disk_guid = uuid.UUID(int=int(os.getenv('SOURCE_DATE_EPOCH')))
-            else:
-                self.disk_guid = uuid.uuid4()
-
             logger.debug("Set disk guid %s", self.disk_guid)
             sfdisk_cmd = "sfdisk --sector-size %s --disk-id %s %s" % \
                         (self.sector_size, self.path, self.disk_guid)