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