X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/5b801534a4277ae5730a12736cd5ff5242aa62d9..bf3a282cb4ca85db501f8a925cd604a1600f25cc:/image_creator/gpt.py?ds=sidebyside

diff --git a/image_creator/gpt.py b/image_creator/gpt.py
index f6f1cc2..968a918 100644
--- a/image_creator/gpt.py
+++ b/image_creator/gpt.py
@@ -42,9 +42,11 @@ BLOCKSIZE = 512
 class MBR(object):
     """Represents a Master Boot Record."""
     class Partition(object):
+        """Represents a partition entry in MBR"""
         format = "<B3sB3sLL"
 
         def __init__(self, raw_part):
+            """Create a Partition instance"""
             (
                 self.status,
                 self.start,
@@ -55,14 +57,14 @@ class MBR(object):
             ) = struct.unpack(self.format, raw_part)
 
         def pack(self):
+            """Pack the partition values into a binary string"""
             return struct.pack(self.format,
-                self.status,
-                self.start,
-                self.type,
-                self.end,
-                self.first_sector,
-                self.sector_count
-            )
+                               self.status,
+                               self.start,
+                               self.type,
+                               self.end,
+                               self.first_sector,
+                               self.sector_count)
 
         @staticmethod
         def size():
@@ -73,7 +75,7 @@ class MBR(object):
             start = self.unpack_chs(self.start)
             end = self.unpack_chs(self.end)
             return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
-                self.first_sector, self.sector_count)
+                                          self.first_sector, self.sector_count)
 
         def unpack_chs(self, chs):
             """Unpacks a CHS address string to a tuple."""
@@ -113,13 +115,14 @@ class MBR(object):
     510     2               MBR signature
     """
     def __init__(self, block):
+        """Create an MBR instance"""
         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.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):
@@ -127,26 +130,25 @@ class MBR(object):
 
     @staticmethod
     def size():
-        """Returns the size of a Master Boot Record."""
+        """Return the size of a Master Boot Record."""
         return struct.calcsize(MBR.format)
 
     def pack(self):
-        """Packs an MBR to a binary string."""
+        """Pack an MBR to a binary string."""
         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
-        )
+                           self.code_area,
+                           self.part[0].pack(),
+                           self.part[1].pack(),
+                           self.part[2].pack(),
+                           self.part[3].pack(),
+                           self.signature)
 
     def __str__(self):
         ret = ""
         for i in range(4):
             ret += "Partition %d: %s\n" % (i, self.part[i])
-        ret += "Signature: %s %s\n" % (
-                    hex(ord(self.signature[0])), hex(ord(self.signature[1])))
+        ret += "Signature: %s %s\n" % (hex(ord(self.signature[0])),
+                                       hex(ord(self.signature[1])))
         return ret
 
 
@@ -176,60 +178,62 @@ class GPTPartitionTable(object):
         """
 
         def __init__(self, block):
-            self.signature, \
-            self.revision, \
-            self.hdr_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)
+            """Create a GPTHeader instance"""
+            (self.signature,
+             self.revision,
+             self.hdr_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):
             """Packs a GPT Header to a binary string."""
             return struct.pack(self.format,
-                self.signature, \
-                self.revision, \
-                self.hdr_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
-            )
+                               self.signature,
+                               self.revision,
+                               self.hdr_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)
 
         @staticmethod
         def size():
-            """Returns the size of a GPT Header."""
+            """Return the size of a GPT Header."""
             return struct.calcsize(GPTPartitionTable.GPTHeader.format)
 
         def __str__(self):
-            return \
-            "Signature: %s\n" % self.signature + \
-            "Revision: %r\n" % self.revision + \
-            "Header Size: %d\n" % self.hdr_size + \
-            "CRC32: %d\n" % self.header_crc32 + \
-            "Current LBA: %d\n" % self.current_lba + \
-            "Backup LBA: %d\n" % self.backup_lba + \
-            "First Usable LBA: %d\n" % self.first_usable_lba + \
-            "Last Usable LBA: %d\n" % self.last_usable_lba + \
-            "Disk GUID: %s\n" % uuid.UUID(bytes=self.uuid) + \
-            "Partition entries starting LBA: %d\n" % self.part_entry_start + \
-            "Number of Partition entries: %d\n" % self.part_count + \
-            "Size of a partition entry: %d\n" % self.part_entry_size + \
-            "CRC32 of partition array: %s\n" % self.part_crc32
+            """Print a GPTHeader"""
+            return "Signature: %s\n" % self.signature + \
+                   "Revision: %r\n" % self.revision + \
+                   "Header Size: %d\n" % self.hdr_size + \
+                   "CRC32: %d\n" % self.header_crc32 + \
+                   "Current LBA: %d\n" % self.current_lba + \
+                   "Backup LBA: %d\n" % self.backup_lba + \
+                   "First Usable LBA: %d\n" % self.first_usable_lba + \
+                   "Last Usable LBA: %d\n" % self.last_usable_lba + \
+                   "Disk GUID: %s\n" % uuid.UUID(bytes=self.uuid) + \
+                   "Partition entries starting LBA: %d\n" % \
+                   self.part_entry_start + \
+                   "Number of Partition entries: %d\n" % self.part_count + \
+                   "Size of a partition entry: %d\n" % self.part_entry_size + \
+                   "CRC32 of partition array: %s\n" % self.part_crc32
 
     def __init__(self, disk):
+        """Create a GPTPartitionTable instance"""
         self.disk = disk
         with open(disk, "rb") as d:
             # MBR (Logical block address 0)
@@ -243,7 +247,7 @@ class GPTPartitionTable(object):
             # 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.primary.part_entry_size
             self.part_entries = d.read(entries_size)
 
             # Secondary GPT Header (LBA -1)
@@ -252,21 +256,27 @@ class GPTPartitionTable(object):
             self.secondary = self.GPTHeader(raw_header)
 
     def size(self):
-        """Returns the payload size of GPT partitioned device."""
+        """Return the payload size of GPT partitioned device."""
         return (self.primary.backup_lba + 1) * BLOCKSIZE
 
-    def shrink(self, size):
+    def shrink(self, size, old_size):
         """Move the secondary GPT Header entries to the address specified by
         size parameter.
         """
-        if size == self.size():
-            return size
 
-        assert size < self.size()
+        # Most partition manipulation programs leave 2048 sector after the last
+        # partition
+        aligned = size + 2048 * BLOCKSIZE
+
+        # new_size is at least: size + Partition Entries + Secondary GPT Header
+        new_size = aligned if aligned <= old_size else \
+            size + len(self.part_entries) + BLOCKSIZE
+
+        assert new_size <= old_size, "The secodary GPT fits in the device"
+
+        if new_size == self.size():
+            return new_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
@@ -278,7 +288,7 @@ class GPTPartitionTable(object):
         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
+            binascii.crc32(self.primary.pack()) & 0xffffffff
 
         # Fix Secondary header
         self.secondary.header_crc32 = 0
@@ -286,13 +296,14 @@ class GPTPartitionTable(object):
         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
+            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(self.primary.pack())
             d.write('\x00' * (BLOCKSIZE - self.primary.size()))
+            d.write(self.part_entries)
             d.seek(self.secondary.part_entry_start * BLOCKSIZE)
             d.write(self.part_entries)
             d.seek(self.primary.backup_lba * BLOCKSIZE)