code.grnet.gr Git - snf-image-creator/blob - image_creator/gpt.py

   1 # Copyright 2012 GRNET S.A. All rights reserved.
   2 #
   3 # Redistribution and use in source and binary forms, with or
   4 # without modification, are permitted provided that the following
   5 # conditions are met:
   6 #
   7 #   1. Redistributions of source code must retain the above
   8 #      copyright notice, this list of conditions and the following
   9 #      disclaimer.
  10 #
  11 #   2. Redistributions in binary form must reproduce the above
  12 #      copyright notice, this list of conditions and the following
  13 #      disclaimer in the documentation and/or other materials
  14 #      provided with the distribution.
  15 #
  16 # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  17 # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18 # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  19 # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  20 # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  21 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  22 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  23 # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  24 # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  25 # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  26 # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  27 # POSSIBILITY OF SUCH DAMAGE.
  28 #
  29 # The views and conclusions contained in the software and
  30 # documentation are those of the authors and should not be
  31 # interpreted as representing official policies, either expressed
  32 # or implied, of GRNET S.A.
  33
  34 import struct
  35 import sys
  36 import uuid
  37 import binascii
  38
  39 BLOCKSIZE = 512
  40
  41
  42 class MBR(object):
  43     """Represents a Master Boot Record."""
  44     class Partition(object):
  45         format = "<B3sB3sLL"
  46
  47         def __init__(self, raw_part):
  48             (
  49                 self.status,
  50                 self.start,
  51                 self.type,
  52                 self.end,
  53                 self.first_sector,
  54                 self.sector_count
  55             ) = struct.unpack(self.format, raw_part)
  56
  57         def pack(self):
  58             return struct.pack(self.format,
  59                 self.status,
  60                 self.start,
  61                 self.type,
  62                 self.end,
  63                 self.first_sector,
  64                 self.sector_count
  65             )
  66
  67         @staticmethod
  68         def size():
  69             """Returns the size of an MBR partition entry"""
  70             return struct.calcsize(MBR.Partition.format)
  71
  72         def __str__(self):
  73             start = self.unpack_chs(self.start)
  74             end = self.unpack_chs(self.end)
  75             return "%d %s %d %s %d %d" % (self.status, start, self.type, end,
  76                 self.first_sector, self.sector_count)
  77
  78         def unpack_chs(self, chs):
  79             """Unpacks a CHS address string to a tuple."""
  80
  81             assert len(chs) == 3
  82
  83             head = struct.unpack('<B', chs[0])[0]
  84             sector = struct.unpack('<B', chs[1])[0] & 0x3f
  85             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
  86                 struct.unpack('<B', chs[2])[0]
  87
  88             return (cylinder, head, sector)
  89
  90         def pack_chs(self, cylinder, head, sector):
  91             """Packs a CHS tuple to an address string."""
  92
  93             assert 1 <= sector <= 63
  94             assert 0 <= cylinder <= 1023
  95             assert 0 <= head <= 255
  96
  97             byte0 = head
  98             byte1 = (cylinder >> 2) & 0xC0 | sector
  99             byte2 = cylinder & 0xff
 100
 101             return struct.pack('<BBB', byte0, byte1, byte2)
 102
 103     format = "<444s2x16s16s16s16s2s"
 104     """
 105     Offset  Length          Contents
 106     0       440(max. 446)   code area
 107     440     2(optional)     disk signature
 108     444     2               Usually nulls
 109     446     16              Partition 0
 110     462     16              Partition 1
 111     478     16              Partition 2
 112     494     16              Partition 3
 113     510     2               MBR signature
 114     """
 115     def __init__(self, block):
 116         raw_part = {}
 117         self.code_area, \
 118         raw_part[0], \
 119         raw_part[1], \
 120         raw_part[2], \
 121         raw_part[3], \
 122         self.signature = struct.unpack(self.format, block)
 123
 124         self.part = {}
 125         for i in range(4):
 126             self.part[i] = self.Partition(raw_part[i])
 127
 128     @staticmethod
 129     def size():
 130         """Returns the size of a Master Boot Record."""
 131         return struct.calcsize(MBR.format)
 132
 133     def pack(self):
 134         """Packs an MBR to a binary string."""
 135         return struct.pack(self.format,
 136             self.code_area,
 137             self.part[0].pack(),
 138             self.part[1].pack(),
 139             self.part[2].pack(),
 140             self.part[3].pack(),
 141             self.signature
 142         )
 143
 144     def __str__(self):
 145         ret = ""
 146         for i in range(4):
 147             ret += "Partition %d: %s\n" % (i, self.part[i])
 148         ret += "Signature: %s %s\n" % (
 149                     hex(ord(self.signature[0])), hex(ord(self.signature[1])))
 150         return ret
 151
 152
 153 class GPTPartitionTable(object):
 154     """Represents a GUID Partition Table."""
 155     class GPTHeader(object):
 156         """Represents a GPT Header of a GUID Partition Table."""
 157         format = "<8s4sII4xQQQQ16sQIII"
 158         """
 159         Offset  Length          Contents
 160         0       8 bytes         Signature
 161         8       4 bytes         Revision
 162         12      4 bytes         Header size in little endian
 163         16      4 bytes         CRC32 of header
 164         20      4 bytes         Reserved; must be zero
 165         24      8 bytes         Current LBA
 166         32      8 bytes         Backup LBA
 167         40      8 bytes         First usable LBA for partitions
 168         48      8 bytes         Last usable LBA
 169         56      16 bytes        Disk GUID
 170         72      8 bytes         Partition entries starting LBA
 171         80      4 bytes         Number of partition entries
 172         84      4 bytes         Size of a partition entry
 173         88      4 bytes         CRC32 of partition array
 174         92      *               Reserved; must be zeroes
 175         LBA    size            Total
 176         """
 177
 178         def __init__(self, block):
 179             self.signature, \
 180             self.revision, \
 181             self.hdr_size, \
 182             self.header_crc32, \
 183             self.current_lba, \
 184             self.backup_lba, \
 185             self.first_usable_lba, \
 186             self.last_usable_lba, \
 187             self.uuid, \
 188             self.part_entry_start, \
 189             self.part_count, \
 190             self.part_entry_size, \
 191             self.part_crc32 = struct.unpack(self.format, block)
 192
 193         def pack(self):
 194             """Packs a GPT Header to a binary string."""
 195             return struct.pack(self.format,
 196                 self.signature, \
 197                 self.revision, \
 198                 self.hdr_size, \
 199                 self.header_crc32, \
 200                 self.current_lba, \
 201                 self.backup_lba, \
 202                 self.first_usable_lba, \
 203                 self.last_usable_lba, \
 204                 self.uuid, \
 205                 self.part_entry_start, \
 206                 self.part_count, \
 207                 self.part_entry_size, \
 208                 self.part_crc32
 209             )
 210
 211         @staticmethod
 212         def size():
 213             """Returns the size of a GPT Header."""
 214             return struct.calcsize(GPTPartitionTable.GPTHeader.format)
 215
 216         def __str__(self):
 217             return \
 218             "Signature: %s\n" % self.signature + \
 219             "Revision: %r\n" % self.revision + \
 220             "Header Size: %d\n" % self.hdr_size + \
 221             "CRC32: %d\n" % self.header_crc32 + \
 222             "Current LBA: %d\n" % self.current_lba + \
 223             "Backup LBA: %d\n" % self.backup_lba + \
 224             "First Usable LBA: %d\n" % self.first_usable_lba + \
 225             "Last Usable LBA: %d\n" % self.last_usable_lba + \
 226             "Disk GUID: %s\n" % uuid.UUID(bytes=self.uuid) + \
 227             "Partition entries starting LBA: %d\n" % self.part_entry_start + \
 228             "Number of Partition entries: %d\n" % self.part_count + \
 229             "Size of a partition entry: %d\n" % self.part_entry_size + \
 230             "CRC32 of partition array: %s\n" % self.part_crc32
 231
 232     def __init__(self, disk):
 233         self.disk = disk
 234         with open(disk, "rb") as d:
 235             # MBR (Logical block address 0)
 236             lba0 = d.read(BLOCKSIZE)
 237             self.mbr = MBR(lba0)
 238
 239             # Primary GPT Header (LBA 1)
 240             raw_header = d.read(self.GPTHeader.size())
 241             self.primary = self.GPTHeader(raw_header)
 242
 243             # Partition entries (LBA 2...34)
 244             d.seek(self.primary.part_entry_start * BLOCKSIZE)
 245             entries_size = self.primary.part_count * \
 246                                                 self.primary.part_entry_size
 247             self.part_entries = d.read(entries_size)
 248
 249             # Secondary GPT Header (LBA -1)
 250             d.seek(self.primary.backup_lba * BLOCKSIZE)
 251             raw_header = d.read(self.GPTHeader.size())
 252             self.secondary = self.GPTHeader(raw_header)
 253
 254     def size(self):
 255         """Returns the payload size of GPT partitioned device."""
 256         return (self.primary.backup_lba + 1) * BLOCKSIZE
 257
 258     def shrink(self, size):
 259         """Move the secondary GPT Header entries to the address specified by
 260         size parameter.
 261         """
 262         if size == self.size():
 263             return size
 264
 265         assert size < self.size()
 266
 267         # new_size = size + Partition Entries + Secondary GPT Header
 268         new_size = size + len(self.part_entries) + BLOCKSIZE
 269         new_size = ((new_size + 4095) // 4096) * 4096  # align to 4K
 270         lba_count = new_size // BLOCKSIZE
 271
 272         # Correct MBR
 273         #TODO: Check if the partition tables is hybrid
 274         self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
 275
 276         # Fix Primary header
 277         self.primary.header_crc32 = 0
 278         self.primary.backup_lba = lba_count - 1  # LBA-1
 279         self.primary.last_usable_lba = lba_count - 34  # LBA-34
 280         self.primary.header_crc32 = \
 281                             binascii.crc32(self.primary.pack()) & 0xffffffff
 282
 283         # Fix Secondary header
 284         self.secondary.header_crc32 = 0
 285         self.secondary.current_lba = self.primary.backup_lba
 286         self.secondary.last_usable_lba = lba_count - 34  # LBA-34
 287         self.secondary.part_entry_start = lba_count - 33  # LBA-33
 288         self.secondary.header_crc32 = \
 289                             binascii.crc32(self.secondary.pack()) & 0xffffffff
 290
 291         # Copy the new partition table back to the device
 292         with open(self.disk, "wb") as d:
 293             d.write(self.mbr.pack())
 294             d.write(self.primary.pack())
 295             d.write('\x00' * (BLOCKSIZE - self.primary.size()))
 296             d.seek(self.secondary.part_entry_start * BLOCKSIZE)
 297             d.write(self.part_entries)
 298             d.seek(self.primary.backup_lba * BLOCKSIZE)
 299             d.write(self.secondary.pack())
 300             d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
 301
 302         return new_size
 303
 304 if __name__ == '__main__':
 305     ptable = GPTPartitionTable(sys.argv[1])
 306
 307     print "MBR:\n%s" % ptable.mbr
 308     print "Primary partition table:\n%s" % ptable.primary
 309     print "Secondary partition table:\n%s" % ptable.secondary
 310
 311 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :