Cleanup code & add comments in gpt module

image_creator/gpt.py

@@ -40,12 +40,13 @@ BLOCKSIZE = 512
 
 
 class MBR(object):
+    """Represents a Master Boot Record."""
     class Partition(object):
         format = "<B3sB3sLL"
 
         def __init__(self, raw_part):
-
-            (   self.status,
+            (
+                self.status,
                 self.start,
                 self.type,
                 self.end,
@@ -63,13 +64,19 @@ class MBR(object):
                 self.sector_count
             )
 
-        def show(self):
+        @staticmethod
+        def size():
+            """Returns the size of an MBR partition entry"""
+            return struct.calcsize(MBR.Partition.format)
+
+        def __str__(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,
+            return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
                 self.first_sector, self.sector_count)
 
         def unpack_chs(self, chs):
+            """Unpacks a CHS address string to a tuple."""
 
             assert len(chs) == 3
 
@@ -81,6 +88,7 @@ class MBR(object):
             return (cylinder, head, sector)
 
         def pack_chs(self, cylinder, head, sector):
+            """Packs a CHS tuple to an address string."""
 
             assert 1 <= sector <= 63
             assert 0 <= cylinder <= 1023
@@ -117,7 +125,13 @@ class MBR(object):
         for i in range(4):
             self.part[i] = self.Partition(raw_part[i])
 
+    @staticmethod
+    def size():
+        """Returns the size of a Master Boot Record."""
+        return struct.calcsize(MBR.format)
+
     def pack(self):
+        """Packs an MBR to a binary string."""
         return struct.pack(self.format,
             self.code_area,
             self.part[0].pack(),
@@ -127,14 +141,19 @@ class MBR(object):
             self.signature
         )
 
-    def show(self):
+    def __str__(self):
+        ret = ""
         for i in range(4):
-            print "Part %d: " % i,
-            self.part[i].show()
+            ret += "Partition %d: %s\n" % (i, self.part[i])
+        ret += "Signature: %s %s\n" % (
+                    hex(ord(self.signature[0])), hex(ord(self.signature[1])))
+        return ret
 
 
 class GPTPartitionTable(object):
+    """Represents a GUID Partition Table."""
     class GPTHeader(object):
+        """Represents a GPT Header of a GUID Partition Table."""
         format = "<8s4sII4xQQQQ16sQIII"
         """
         Offset	Length 	        Contents
@@ -159,7 +178,7 @@ class GPTPartitionTable(object):
         def __init__(self, block):
             self.signature, \
             self.revision, \
-            self.size, \
+            self.hdr_size, \
             self.header_crc32, \
             self.current_lba, \
             self.backup_lba, \
@@ -172,10 +191,11 @@ class GPTPartitionTable(object):
             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.size, \
+                self.hdr_size, \
                 self.header_crc32, \
                 self.current_lba, \
                 self.backup_lba, \
@@ -188,20 +208,26 @@ class GPTPartitionTable(object):
                 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
+        @staticmethod
+        def size():
+            """Returns 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
 
     def __init__(self, disk):
         self.disk = disk
@@ -209,24 +235,30 @@ class GPTPartitionTable(object):
             #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])
+            raw_header = d.read(self.GPTHeader.size())
+            self.primary = self.GPTHeader(raw_header)
+
             # 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])
+            raw_header = d.read(self.GPTHeader.size())
+            self.secondary = self.GPTHeader(raw_header)
 
     def size(self):
-        return (self.primary.backup_lba + 1) * BLOCKSIZE 
+        """Returns the payload size of GPT partitioned device."""
+        return (self.primary.backup_lba + 1) * BLOCKSIZE
 
     def shrink(self, size):
-
+        """Move the secondary GPT Header entries to the address specified by
+        size parameter.
+        """
         if size == self.size():
             return size
 
@@ -238,17 +270,17 @@ class GPTPartitionTable(object):
         lba_count = new_size // BLOCKSIZE
 
         # Correct MBR
-        #TODO: Check for hybrid partition tables
+        #TODO: Check if the partition tables is hybrid
         self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
 
-        # Correct Primary header
+        # Fix 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
+        # Fix Secondary header
         self.secondary.header_crc32 = 0
         self.secondary.current_lba = self.primary.backup_lba
         self.secondary.last_usable_lba = lba_count - 34  # LBA-34
@@ -259,28 +291,21 @@ class GPTPartitionTable(object):
         # 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.write('\x00' * (BLOCKSIZE - self.primary.size()))
             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())
+            d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
 
         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()
+    print "MBR:\n%s" % ptable.mbr
+    print "Primary partition table:\n%s" % ptable.primary
+    print "Secondary partition table:\n%s" % ptable.secondary
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :