Skip to content
Snippets Groups Projects
Commit c0240ac1 authored by Nikos Skalkotos's avatar Nikos Skalkotos
Browse files

Cleanup code & add comments in gpt module

parent 331aa0ec
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Showing
with 69 additions and 44 deletions
image_creator/gpt.py
+ 69
− 44
View file @ c0240ac1
......@@ -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 :
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment