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     class Partition(object):
  44         format = "<B3sB3sLL"
  45
  46         def __init__(self, raw_part):
  47
  48             (   self.status,
  49                 self.start,
  50                 self.type,
  51                 self.end,
  52                 self.first_sector,
  53                 self.sector_count
  54             ) = struct.unpack(self.format, raw_part)
  55
  56         def pack(self):
  57             return struct.pack(self.format,
  58                 self.status,
  59                 self.start,
  60                 self.type,
  61                 self.end,
  62                 self.first_sector,
  63                 self.sector_count
  64             )
  65
  66         def show(self):
  67             start = self.unpack_chs(self.start)
  68             end = self.unpack_chs(self.end)
  69             print "%d %s %d %s %d %d" % (self.status, start, self.type, end,
  70                 self.first_sector, self.sector_count)
  71
  72         def unpack_chs(self, chs):
  73
  74             assert len(chs) == 3
  75
  76             head = struct.unpack('<B', chs[0])[0]
  77             sector = struct.unpack('<B', chs[1])[0] & 0x3f
  78             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
  79                 struct.unpack('<B', chs[2])[0]
  80
  81             return (cylinder, head, sector)
  82
  83         def pack_chs(self, cylinder, head, sector):
  84
  85             assert 1 <= sector <= 63
  86             assert 0 <= cylinder <= 1023
  87             assert 0 <= head <= 255
  88
  89             byte0 = head
  90             byte1 = (cylinder >> 2) & 0xC0 | sector
  91             byte2 = cylinder & 0xff
  92
  93             return struct.pack('<BBB', byte0, byte1, byte2)
  94
  95     format = "<444s2x16s16s16s16s2s"
  96     """
  97     Offset  Length          Contents
  98     0       440(max. 446)   code area
  99     440     2(optional)     disk signature
 100     444     2               Usually nulls
 101     446     16              Partition 0
 102     462     16              Partition 1
 103     478     16              Partition 2
 104     494     16              Partition 3
 105     510     2               MBR signature
 106     """
 107     def __init__(self, block):
 108         raw_part = {}
 109         self.code_area, \
 110         raw_part[0], \
 111         raw_part[1], \
 112         raw_part[2], \
 113         raw_part[3], \
 114         self.signature = struct.unpack(self.format, block)
 115
 116         self.part = {}
 117         for i in range(4):
 118             self.part[i] = self.Partition(raw_part[i])
 119
 120     def pack(self):
 121         return struct.pack(self.format,
 122             self.code_area,
 123             self.part[0].pack(),
 124             self.part[1].pack(),
 125             self.part[2].pack(),
 126             self.part[3].pack(),
 127             self.signature
 128         )
 129
 130     def show(self):
 131         for i in range(4):
 132             print "Part %d: " % i,
 133             self.part[i].show()
 134
 135
 136 class GPTPartitionTable(object):
 137     class GPTHeader(object):
 138         format = "<8s4sII4xQQQQ16sQIII"
 139         """
 140         Offset  Length          Contents
 141         0       8 bytes         Signature
 142         8       4 bytes         Revision
 143         12      4 bytes         Header size in little endian
 144         16      4 bytes         CRC32 of header
 145         20      4 bytes         Reserved; must be zero
 146         24      8 bytes         Current LBA
 147         32      8 bytes         Backup LBA
 148         40      8 bytes         First usable LBA for partitions
 149         48      8 bytes         Last usable LBA
 150         56      16 bytes        Disk GUID
 151         72      8 bytes         Partition entries starting LBA
 152         80      4 bytes         Number of partition entries
 153         84      4 bytes         Size of a partition entry
 154         88      4 bytes         CRC32 of partition array
 155         92      *               Reserved; must be zeroes
 156         LBA    size            Total
 157         """
 158
 159         def __init__(self, block):
 160             self.signature, \
 161             self.revision, \
 162             self.size, \
 163             self.header_crc32, \
 164             self.current_lba, \
 165             self.backup_lba, \
 166             self.first_usable_lba, \
 167             self.last_usable_lba, \
 168             self.uuid, \
 169             self.part_entry_start, \
 170             self.part_count, \
 171             self.part_entry_size, \
 172             self.part_crc32 = struct.unpack(self.format, block)
 173
 174         def pack(self):
 175             return struct.pack(self.format,
 176                 self.signature, \
 177                 self.revision, \
 178                 self.size, \
 179                 self.header_crc32, \
 180                 self.current_lba, \
 181                 self.backup_lba, \
 182                 self.first_usable_lba, \
 183                 self.last_usable_lba, \
 184                 self.uuid, \
 185                 self.part_entry_start, \
 186                 self.part_count, \
 187                 self.part_entry_size, \
 188                 self.part_crc32
 189             )
 190
 191         def show(self):
 192             print "Signature: %s" % self.signature
 193             print "Revision: %r" % self.revision
 194             print "Header Size: %d" % self.size
 195             print "CRC32: %d" % self.header_crc32
 196             print "Current LBA: %d" % self.current_lba
 197             print "Backup LBA: %d" % self.backup_lba
 198             print "First Usable LBA: %d" % self.first_usable_lba
 199             print "Last Usable LBA: %d" % self.last_usable_lba
 200             print "Disk GUID: %s" % uuid.UUID(bytes=self.uuid)
 201             print "Partition entries starting LBA: %d" % self.part_entry_start
 202             print "Number of Partition entries: %d" % self.part_count
 203             print "Size of a partition entry: %d" % self.part_entry_size
 204             print "CRC32 of partition array: %s" % self.part_crc32
 205
 206     def __init__(self, disk):
 207         self.disk = disk
 208         with open(disk, "rb") as d:
 209             #MBR (Logical block address 0)
 210             lba0 = d.read(BLOCKSIZE)
 211             self.mbr = MBR(lba0)
 212             # Primary GPT Header (LBA 1)
 213             lba1 = d.read(BLOCKSIZE)
 214             self.primary = self.GPTHeader(lba1[:92])
 215             # Partition entries (LBA 2...34)
 216             d.seek(self.primary.part_entry_start * BLOCKSIZE)
 217             entries_size = self.primary.part_count * \
 218                                                 self.primary.part_entry_size
 219             self.part_entries = d.read(entries_size)
 220             # Secondary GPT Header (LBA -1)
 221             d.seek(self.primary.backup_lba * BLOCKSIZE)
 222             lba_1 = d.read(BLOCKSIZE)
 223             self.secondary = self.GPTHeader(lba_1[:92])
 224
 225     def size(self):
 226         return (self.primary.backup_lba + 1) * BLOCKSIZE
 227
 228     def shrink(self, size):
 229
 230         if size == self.size():
 231             return size
 232
 233         assert size < self.size()
 234
 235         # new_size = size + Partition Entries + Secondary GPT Header
 236         new_size = size + len(self.part_entries) + BLOCKSIZE
 237         new_size = ((new_size + 4095) // 4096) * 4096  # align to 4K
 238         lba_count = new_size // BLOCKSIZE
 239
 240         # Correct MBR
 241         #TODO: Check for hybrid partition tables
 242         self.mbr.part[0].sector_count = (new_size // BLOCKSIZE) - 1
 243
 244         # Correct Primary header
 245         self.primary.header_crc32 = 0
 246         self.primary.backup_lba = lba_count - 1  # LBA-1
 247         self.primary.last_usable_lba = lba_count - 34  # LBA-34
 248         self.primary.header_crc32 = \
 249                             binascii.crc32(self.primary.pack()) & 0xffffffff
 250
 251         # Correct Secondary header entries
 252         self.secondary.header_crc32 = 0
 253         self.secondary.current_lba = self.primary.backup_lba
 254         self.secondary.last_usable_lba = lba_count - 34  # LBA-34
 255         self.secondary.part_entry_start = lba_count - 33  # LBA-33
 256         self.secondary.header_crc32 = \
 257                             binascii.crc32(self.secondary.pack()) & 0xffffffff
 258
 259         # Copy the new partition table back to the device
 260         with open(self.disk, "wb") as d:
 261             d.write(self.mbr.pack())
 262             d.write(struct.pack("%ss" % BLOCKSIZE, '\x00' * BLOCKSIZE))
 263             d.seek(BLOCKSIZE)
 264             d.write(self.primary.pack())
 265             d.seek(self.secondary.part_entry_start * BLOCKSIZE)
 266             d.write(self.part_entries)
 267             d.seek(self.primary.backup_lba * BLOCKSIZE)
 268             d.write(struct.pack("%ss" % BLOCKSIZE, '\x00' * BLOCKSIZE))
 269             d.seek(self.primary.backup_lba * BLOCKSIZE)
 270             d.write(self.secondary.pack())
 271
 272         return new_size
 273
 274 if __name__ == '__main__':
 275     ptable = GPTPartitionTable(sys.argv[1])
 276
 277     print "MBR:"
 278     ptable.mbr.show()
 279     print
 280     print "Primary partition table:"
 281     ptable.primary.show()
 282     print
 283     print "Secondary partition table:"
 284     ptable.secondary.show()
 285
 286 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :