diff --git a/image_creator/gpt.py b/image_creator/gpt.py index f85305848cf4025923acf6e2a736b10f4c1f3467..1a7ed04c99810d6d5ee86168c0e5bb06bca39918 100644 --- a/image_creator/gpt.py +++ b/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 :