From: Nikos Skalkotos <skalkoto@grnet.gr>
Date: Thu, 26 Apr 2012 16:55:53 +0000 (+0300)
Subject: Add support for GUID Partition Tables
X-Git-Tag: v0.1~98
X-Git-Url: https://code.grnet.gr/git/snf-image-creator/commitdiff_plain/331aa0ec4a14fb2484e819961222297cb022e190?ds=sidebyside

Add support for GUID Partition Tables
---

diff --git a/image_creator/disk.py b/image_creator/disk.py
index 287921a..79ce422 100644
--- a/image_creator/disk.py
+++ b/image_creator/disk.py
@@ -33,7 +33,7 @@
 
 from image_creator.util import get_command
 from image_creator.util import warn, progress, success, output, FatalError
-
+from image_creator.gpt import GPTPartitionTable
 import stat
 import os
 import tempfile
@@ -200,6 +200,9 @@ class DiskDevice(object):
             raise FatalError("Multiple operating systems found."
                             "We only support images with one filesystem.")
         self.root = roots[0]
+        self.gdev = self.g.part_to_dev(self.root)
+        self.parttype = self.g.part_get_parttype(self.gdev)
+
         self.ostype = self.g.inspect_get_type(self.root)
         self.distro = self.g.inspect_get_distro(self.root)
         success('found a(n) %s system' % self.distro)
@@ -256,23 +259,21 @@ class DiskDevice(object):
         """
         output("Shrinking image (this may take a while)...", False)
 
-        dev = self.g.part_to_dev(self.root)
-        parttype = self.g.part_get_parttype(dev)
-        if parttype != 'msdos':
+        if self.parttype not in 'msdos' 'gpt':
             raise FatalError("You have a %s partition table. "
-                "Only msdos partitions are supported" % parttype)
+                "Only msdos and gpt partitions are supported" % self.parttype)
 
-        last_partition = self.g.part_list(dev)[-1]
+        last_partition = self.g.part_list(self.gdev)[-1]
 
         if last_partition['part_num'] > 4:
             raise FatalError("This disk contains logical partitions. "
                 "Only primary partitions are supported.")
 
-        part_dev = "%s%d" % (dev, last_partition['part_num'])
+        part_dev = "%s%d" % (self.gdev, last_partition['part_num'])
         fs_type = self.g.vfs_type(part_dev)
         if not re.match("ext[234]", fs_type):
             warn("Don't know how to resize %s partitions." % vfs_type)
-            return
+            return self.size()
 
         self.g.e2fsck_f(part_dev)
         self.g.resize2fs_M(part_dev)
@@ -283,17 +284,22 @@ class DiskDevice(object):
         block_cnt = int(
             filter(lambda x: x[0] == 'Block count', out)[0][1])
 
-        sector_size = self.g.blockdev_getss(dev)
+        sector_size = self.g.blockdev_getss(self.gdev)
 
         start = last_partition['part_start'] / sector_size
         end = start + (block_size * block_cnt) / sector_size - 1
 
-        self.g.part_del(dev, last_partition['part_num'])
-        self.g.part_add(dev, 'p', start, end)
+        self.g.part_del(self.gdev, last_partition['part_num'])
+        self.g.part_add(self.gdev, 'p', start, end)
 
         new_size = (end + 1) * sector_size
         success("new image size is %dMB" %
                             ((new_size + 2 ** 20 - 1) // 2 ** 20))
+
+        if self.parttype == 'gpt':
+            ptable = GPTPartitionTable(self.device)
+            return ptable.shrink(new_size)
+
         return new_size
 
     def size(self):
@@ -302,10 +308,16 @@ class DiskDevice(object):
         The size returned by this method is the size of the space occupied by
         the partitions (including the space before the first partition).
         """
-        dev = self.g.part_to_dev(self.root)
-        last = self.g.part_list(dev)[-1]
 
-        return last['part_end'] + 1
+        if self.parttype == 'msdos':
+            dev = self.g.part_to_dev(self.root)
+            last = self.g.part_list(dev)[-1]
+            return last['part_end'] + 1
+        elif self.parttype == 'gpt':
+            ptable = GPTPartitionTable(self.device)
+            return ptable.size()
+        else:
+            raise FatalError("Unsupported partition table type: %s" % parttype)
 
     def dump(self, outfile):
         """Dumps the content of device into a file.
diff --git a/image_creator/gpt.py b/image_creator/gpt.py
new file mode 100644
index 0000000..f853058
--- /dev/null
+++ b/image_creator/gpt.py
@@ -0,0 +1,286 @@
+# Copyright 2012 GRNET S.A. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or
+# without modification, are permitted provided that the following
+# conditions are met:
+#
+#   1. Redistributions of source code must retain the above
+#      copyright notice, this list of conditions and the following
+#      disclaimer.
+#
+#   2. Redistributions in binary form must reproduce the above
+#      copyright notice, this list of conditions and the following
+#      disclaimer in the documentation and/or other materials
+#      provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
+# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
+# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
+# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+# The views and conclusions contained in the software and
+# documentation are those of the authors and should not be
+# interpreted as representing official policies, either expressed
+# or implied, of GRNET S.A.
+
+import struct
+import sys
+import uuid
+import binascii
+
+BLOCKSIZE = 512
+
+
+class MBR(object):
+    class Partition(object):
+        format = "<B3sB3sLL"
+
+        def __init__(self, raw_part):
+
+            (   self.status,
+                self.start,
+                self.type,
+                self.end,
+                self.first_sector,
+                self.sector_count
+            ) = struct.unpack(self.format, raw_part)
+
+        def pack(self):
+            return struct.pack(self.format,
+                self.status,
+                self.start,
+                self.type,
+                self.end,
+                self.first_sector,
+                self.sector_count
+            )
+
+        def show(self):
+            start = self.unpack_chs(self.start)
+            end = self.unpack_chs(self.end)
+            print "%d %s %d %s %d %d" % (self.status, start, self.type, end,
+                self.first_sector, self.sector_count)
+
+        def unpack_chs(self, chs):
+
+            assert len(chs) == 3
+
+            head = struct.unpack('<B', chs[0])[0]
+            sector = struct.unpack('<B', chs[1])[0] & 0x3f
+            cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
+                struct.unpack('<B', chs[2])[0]
+
+            return (cylinder, head, sector)
+
+        def pack_chs(self, cylinder, head, sector):
+
+            assert 1 <= sector <= 63
+            assert 0 <= cylinder <= 1023
+            assert 0 <= head <= 255
+
+            byte0 = head
+            byte1 = (cylinder >> 2) & 0xC0 | sector
+            byte2 = cylinder & 0xff
+
+            return struct.pack('<BBB', byte0, byte1, byte2)
+
+    format = "<444s2x16s16s16s16s2s"
+    """
+    Offset  Length          Contents
+    0       440(max. 446)   code area
+    440     2(optional)     disk signature
+    444     2               Usually nulls
+    446     16              Partition 0
+    462     16              Partition 1
+    478     16              Partition 2
+    494     16              Partition 3
+    510     2               MBR signature
+    """
+    def __init__(self, block):
+        raw_part = {}
+        self.code_area, \
+        raw_part[0], \
+        raw_part[1], \
+        raw_part[2], \
+        raw_part[3], \
+        self.signature = struct.unpack(self.format, block)
+
+        self.part = {}
+        for i in range(4):
+            self.part[i] = self.Partition(raw_part[i])
+
+    def pack(self):
+        return struct.pack(self.format,
+            self.code_area,
+            self.part[0].pack(),
+            self.part[1].pack(),
+            self.part[2].pack(),
+            self.part[3].pack(),
+            self.signature
+        )
+
+    def show(self):
+        for i in range(4):
+            print "Part %d: " % i,
+            self.part[i].show()
+
+
+class GPTPartitionTable(object):
+    class GPTHeader(object):
+        format = "<8s4sII4xQQQQ16sQIII"
+        """
+        Offset	Length 	        Contents
+        0       8 bytes         Signature
+        8       4 bytes 	Revision
+        12      4 bytes 	Header size in little endian
+        16 	4 bytes 	CRC32 of header
+        20 	4 bytes 	Reserved; must be zero
+        24 	8 bytes 	Current LBA
+        32 	8 bytes 	Backup LBA
+        40 	8 bytes 	First usable LBA for partitions
+        48 	8 bytes 	Last usable LBA
+        56 	16 bytes 	Disk GUID
+        72 	8 bytes 	Partition entries starting LBA
+        80 	4 bytes 	Number of partition entries
+        84 	4 bytes 	Size of a partition entry
+        88 	4 bytes 	CRC32 of partition array
+        92 	* 	        Reserved; must be zeroes
+        LBA    size            Total
+        """
+
+        def __init__(self, block):
+            self.signature, \
+            self.revision, \
+            self.size, \
+            self.header_crc32, \
+            self.current_lba, \
+            self.backup_lba, \
+            self.first_usable_lba, \
+            self.last_usable_lba, \
+            self.uuid, \
+            self.part_entry_start, \
+            self.part_count, \
+            self.part_entry_size, \
+            self.part_crc32 = struct.unpack(self.format, block)
+
+        def pack(self):
+            return struct.pack(self.format,
+                self.signature, \
+                self.revision, \
+                self.size, \
+                self.header_crc32, \
+                self.current_lba, \
+                self.backup_lba, \
+                self.first_usable_lba, \
+                self.last_usable_lba, \
+                self.uuid, \
+                self.part_entry_start, \
+                self.part_count, \
+                self.part_entry_size, \
+                self.part_crc32
+            )
+
+        def show(self):
+            print "Signature: %s" % self.signature
+            print "Revision: %r" % self.revision
+            print "Header Size: %d" % self.size
+            print "CRC32: %d" % self.header_crc32
+            print "Current LBA: %d" % self.current_lba
+            print "Backup LBA: %d" % self.backup_lba
+            print "First Usable LBA: %d" % self.first_usable_lba
+            print "Last Usable LBA: %d" % self.last_usable_lba
+            print "Disk GUID: %s" % uuid.UUID(bytes=self.uuid)
+            print "Partition entries starting LBA: %d" % self.part_entry_start
+            print "Number of Partition entries: %d" % self.part_count
+            print "Size of a partition entry: %d" % self.part_entry_size
+            print "CRC32 of partition array: %s" % self.part_crc32
+
+    def __init__(self, disk):
+        self.disk = disk
+        with open(disk, "rb") as d:
+            #MBR (Logical block address 0)
+            lba0 = d.read(BLOCKSIZE)
+            self.mbr = MBR(lba0)
+            # Primary GPT Header (LBA 1)
+            lba1 = d.read(BLOCKSIZE)
+            self.primary = self.GPTHeader(lba1[:92])
+            # Partition entries (LBA 2...34)
+            d.seek(self.primary.part_entry_start * BLOCKSIZE)
+            entries_size = self.primary.part_count * \
+                                                self.primary.part_entry_size
+            self.part_entries = d.read(entries_size)
+            # Secondary GPT Header (LBA -1)
+            d.seek(self.primary.backup_lba * BLOCKSIZE)
+            lba_1 = d.read(BLOCKSIZE)
+            self.secondary = self.GPTHeader(lba_1[:92])
+
+    def size(self):
+        return (self.primary.backup_lba + 1) * BLOCKSIZE 
+
+    def shrink(self, size):
+
+        if size == self.size():
+            return size
+
+        assert size < self.size()
+
+        # new_size = size + Partition Entries + Secondary GPT Header
+        new_size = size + len(self.part_entries) + BLOCKSIZE
+        new_size = ((new_size + 4095) // 4096) * 4096  # align to 4K
+        lba_count = new_size // BLOCKSIZE
+
+        # Correct MBR
+        #TODO: Check for hybrid partition tables
+        self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
+
+        # Correct Primary header
+        self.primary.header_crc32 = 0
+        self.primary.backup_lba = lba_count - 1  # LBA-1
+        self.primary.last_usable_lba = lba_count - 34  # LBA-34
+        self.primary.header_crc32 = \
+                            binascii.crc32(self.primary.pack()) & 0xffffffff
+
+        # Correct Secondary header entries
+        self.secondary.header_crc32 = 0
+        self.secondary.current_lba = self.primary.backup_lba
+        self.secondary.last_usable_lba = lba_count - 34  # LBA-34
+        self.secondary.part_entry_start = lba_count - 33  # LBA-33
+        self.secondary.header_crc32 = \
+                            binascii.crc32(self.secondary.pack()) & 0xffffffff
+
+        # Copy the new partition table back to the device
+        with open(self.disk, "wb") as d:
+            d.write(self.mbr.pack())
+            d.write(struct.pack("%ss" % BLOCKSIZE, '\x00' * BLOCKSIZE))
+            d.seek(BLOCKSIZE)
+            d.write(self.primary.pack())
+            d.seek(self.secondary.part_entry_start * BLOCKSIZE)
+            d.write(self.part_entries)
+            d.seek(self.primary.backup_lba * BLOCKSIZE)
+            d.write(struct.pack("%ss" % BLOCKSIZE, '\x00' * BLOCKSIZE))
+            d.seek(self.primary.backup_lba * BLOCKSIZE)
+            d.write(self.secondary.pack())
+
+        return new_size
+
+if __name__ == '__main__':
+    ptable = GPTPartitionTable(sys.argv[1])
+
+    print "MBR:"
+    ptable.mbr.show()
+    print
+    print "Primary partition table:"
+    ptable.primary.show()
+    print
+    print "Secondary partition table:"
+    ptable.secondary.show()
+
+# vim: set sta sts=4 shiftwidth=4 sw=4 et ai :
diff --git a/image_creator/util.py b/image_creator/util.py
index efccabb..d3d4cc0 100644
--- a/image_creator/util.py
+++ b/image_creator/util.py
@@ -113,7 +113,7 @@ def md5(filename, size):
     BLOCKSIZE = 2 ** 22  # 4MB
 
     progressbar = progress("Calculating md5sum:", 'mb')
-    progressbar.max = (size // (2 ** 20))
+    progressbar.max = ((size + 2 ** 20 - 1) // (2 ** 20))
     md5 = hashlib.md5()
     with open(filename, "r") as src:
         left = size