mic/utils/gpt_parser.py

   1 #!/usr/bin/python -tt
   2 #
   3 # Copyright (c) 2013 Intel, Inc.
   4 #
   5 # This program is free software; you can redistribute it and/or modify it
   6 # under the terms of the GNU General Public License as published by the Free
   7 # Software Foundation; version 2 of the License
   8 #
   9 # This program is distributed in the hope that it will be useful, but
  10 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  11 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  12 # for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License along
  15 # with this program; if not, write to the Free Software Foundation, Inc., 59
  16 # Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  17
  18 """ This module implements a simple GPT partitions parser which can read the
  19 GPT header and the GPT partition table. """
  20
  21 import struct
  22 import uuid
  23 import binascii
  24 from mic.utils.errors import MountError
  25
  26 _GPT_HEADER_FORMAT = "<8s4sIIIQQQQ16sQIII"
  27 _GPT_HEADER_SIZE = struct.calcsize(_GPT_HEADER_FORMAT)
  28 _GPT_ENTRY_FORMAT = "<16s16sQQQ72s"
  29 _GPT_ENTRY_SIZE = struct.calcsize(_GPT_ENTRY_FORMAT)
  30 _SUPPORTED_GPT_REVISION = '\x00\x00\x01\x00'
  31
  32 def _stringify_uuid(binary_uuid):
  33     """ A small helper function to transform a binary UUID into a string
  34     format. """
  35
  36     uuid_str = str(uuid.UUID(bytes_le = binary_uuid))
  37
  38     return uuid_str.upper()
  39
  40 def _calc_header_crc(raw_hdr):
  41     """ Calculate GPT header CRC32 checksum. The 'raw_hdr' parameter has to
  42     be a list or a tuple containing all the elements of the GPT header in a
  43     "raw" form, meaning that it should simply contain "unpacked" disk data.
  44     """
  45
  46     raw_hdr = list(raw_hdr)
  47     raw_hdr[3] = 0
  48     raw_hdr = struct.pack(_GPT_HEADER_FORMAT, *raw_hdr)
  49
  50     return binascii.crc32(raw_hdr) & 0xFFFFFFFF
  51
  52 def _validate_header(raw_hdr):
  53     """ Validate the GPT header. The 'raw_hdr' parameter has to be a list or a
  54     tuple containing all the elements of the GPT header in a "raw" form,
  55     meaning that it should simply contain "unpacked" disk data. """
  56
  57     # Validate the signature
  58     if raw_hdr[0] != 'EFI PART':
  59         raise MountError("GPT partition table not found")
  60
  61     # Validate the revision
  62     if raw_hdr[1] != _SUPPORTED_GPT_REVISION:
  63         raise MountError("Unsupported GPT revision '%s', supported revision " \
  64                          "is '%s'" % \
  65                           (binascii.hexlify(raw_hdr[1]),
  66                            binascii.hexlify(_SUPPORTED_GPT_REVISION)))
  67
  68     # Validate header size
  69     if raw_hdr[2] != _GPT_HEADER_SIZE:
  70         raise MountError("Bad GPT header size: %d bytes, expected %d" % \
  71                          (raw_hdr[2], _GPT_HEADER_SIZE))
  72
  73     crc = _calc_header_crc(raw_hdr)
  74     if raw_hdr[3] != crc:
  75         raise MountError("GPT header crc mismatch: %#x, should be %#x" % \
  76                          (crc, raw_hdr[3]))
  77
  78 class GptParser:
  79     """ GPT partition table parser. Allows reading the GPT header and the
  80     partition table, as well as modifying the partition table records. """
  81
  82     def __init__(self, disk_path, sector_size = 512):
  83         """ The class constructor which accepts the following parameters:
  84             * disk_path - full path to the disk image or device node
  85             * sector_size - size of a disk sector in bytes """
  86
  87         self.sector_size = sector_size
  88         self.disk_path = disk_path
  89
  90         try:
  91             self._disk_obj = open(disk_path, 'r+b')
  92         except IOError as err:
  93             raise MountError("Cannot open file '%s' for reading GPT " \
  94                              "partitions: %s" % (disk_path, err))
  95
  96     def __del__(self):
  97         """ The class destructor. """
  98
  99         self._disk_obj.close()
 100
 101     def _read_disk(self, offset, size):
 102         """ A helper function which reads 'size' bytes from offset 'offset' of
 103         the disk and checks all the error conditions. """
 104
 105         self._disk_obj.seek(offset)
 106         try:
 107             data = self._disk_obj.read(size)
 108         except IOError as err:
 109             raise MountError("cannot read from '%s': %s" % \
 110                              (self.disk_path, err))
 111
 112         if len(data) != size:
 113             raise MountError("cannot read %d bytes from offset '%d' of '%s', " \
 114                              "read only %d bytes" % \
 115                              (size, offset, self.disk_path, len(data)))
 116
 117         return data
 118
 119     def _write_disk(self, offset, buf):
 120         """ A helper function which writes buffer 'buf' to offset 'offset' of
 121         the disk. This function takes care of unaligned writes and checks all
 122         the error conditions. """
 123
 124         # Since we may be dealing with a block device, we only can write in
 125         # 'self.sector_size' chunks. Find the aligned starting and ending
 126         # disk offsets to read.
 127         start =  (offset / self.sector_size) * self.sector_size
 128         end = ((start + len(buf)) / self.sector_size + 1) * self.sector_size
 129
 130         data = self._read_disk(start, end - start)
 131         off = offset - start
 132         data = data[:off] + buf + data[off + len(buf):]
 133
 134         self._disk_obj.seek(start)
 135         try:
 136             self._disk_obj.write(data)
 137         except IOError as err:
 138             raise MountError("cannot write to '%s': %s" % (self.disk_path, err))
 139
 140     def read_header(self, primary = True):
 141         """ Read and verify the GPT header and return a dictionary containing
 142         the following elements:
 143
 144         'signature'   : header signature
 145         'revision'    : header revision
 146         'hdr_size'    : header size in bytes
 147         'hdr_crc'     : header CRC32
 148         'hdr_lba'     : LBA of this header
 149         'hdr_offs'    : byte disk offset of this header
 150         'backup_lba'  : backup header LBA
 151         'backup_offs' : byte disk offset of backup header
 152         'first_lba'   : first usable LBA for partitions
 153         'first_offs'  : first usable byte disk offset for partitions
 154         'last_lba'    : last usable LBA for partitions
 155         'last_offs'   : last usable byte disk offset for partitions
 156         'disk_uuid'   : UUID of the disk
 157         'ptable_lba'  : starting LBA of array of partition entries
 158         'ptable_offs' : disk byte offset of the start of the partition table
 159         'ptable_size' : partition table size in bytes
 160         'entries_cnt' : number of available partition table entries
 161         'entry_size'  : size of a single partition entry
 162         'ptable_crc'  : CRC32 of the partition table
 163         'primary'     : a boolean, if 'True', this is the primary GPT header,
 164                         if 'False' - the secondary
 165         'primary_str' : contains string "primary" if this is the primary GPT
 166                         header, and "backup" otherwise
 167
 168         This dictionary corresponds to the GPT header format. Please, see the
 169         UEFI standard for the description of these fields.
 170
 171         If the 'primary' parameter is 'True', the primary GPT header is read,
 172         otherwise the backup GPT header is read instead. """
 173
 174         # Read and validate the primary GPT header
 175         raw_hdr = self._read_disk(self.sector_size, _GPT_HEADER_SIZE)
 176         raw_hdr = struct.unpack(_GPT_HEADER_FORMAT, raw_hdr)
 177         _validate_header(raw_hdr)
 178         primary_str = "primary"
 179
 180         if not primary:
 181             # Read and validate the backup GPT header
 182             raw_hdr = self._read_disk(raw_hdr[6] * self.sector_size, _GPT_HEADER_SIZE)
 183             raw_hdr = struct.unpack(_GPT_HEADER_FORMAT, raw_hdr)
 184             _validate_header(raw_hdr)
 185             primary_str = "backup"
 186
 187         return { 'signature'   : raw_hdr[0],
 188                  'revision'    : raw_hdr[1],
 189                  'hdr_size'    : raw_hdr[2],
 190                  'hdr_crc'     : raw_hdr[3],
 191                  'hdr_lba'     : raw_hdr[5],
 192                  'hdr_offs'    : raw_hdr[5] * self.sector_size,
 193                  'backup_lba'  : raw_hdr[6],
 194                  'backup_offs' : raw_hdr[6] * self.sector_size,
 195                  'first_lba'   : raw_hdr[7],
 196                  'first_offs'  : raw_hdr[7] * self.sector_size,
 197                  'last_lba'    : raw_hdr[8],
 198                  'last_offs'   : raw_hdr[8] * self.sector_size,
 199                  'disk_uuid'   :_stringify_uuid(raw_hdr[9]),
 200                  'ptable_lba'  : raw_hdr[10],
 201                  'ptable_offs' : raw_hdr[10] * self.sector_size,
 202                  'ptable_size' : raw_hdr[11] * raw_hdr[12],
 203                  'entries_cnt' : raw_hdr[11],
 204                  'entry_size'  : raw_hdr[12],
 205                  'ptable_crc'  : raw_hdr[13],
 206                  'primary'     : primary,
 207                  'primary_str' : primary_str }
 208
 209     def _read_raw_ptable(self, header):
 210         """ Read and validate primary or backup partition table. The 'header'
 211         argument is the GPT header. If it is the primary GPT header, then the
 212         primary partition table is read and validated, otherwise - the backup
 213         one. The 'header' argument is a dictionary which is returned by the
 214         'read_header()' method. """
 215
 216         raw_ptable = self._read_disk(header['ptable_offs'],
 217                                      header['ptable_size'])
 218
 219         crc = binascii.crc32(raw_ptable) & 0xFFFFFFFF
 220         if crc != header['ptable_crc']:
 221             raise MountError("Partition table at LBA %d (%s) is corrupted" % \
 222                              (header['ptable_lba'], header['primary_str']))
 223
 224         return raw_ptable
 225
 226     def get_partitions(self, primary = True):
 227         """ This is a generator which parses the GPT partition table and
 228         generates the following dictionary for each partition:
 229
 230         'index'       : the index of the partition table endry
 231         'offs'        : byte disk offset of the partition table entry
 232         'type_uuid'   : partition type UUID
 233         'part_uuid'   : partition UUID
 234         'first_lba'   : the first LBA
 235         'last_lba'    : the last LBA
 236         'flags'       : attribute flags
 237         'name'        : partition name
 238         'primary'     : a boolean, if 'True', this is the primary partition
 239                         table, if 'False' - the secondary
 240         'primary_str' : contains string "primary" if this is the primary GPT
 241                         header, and "backup" otherwise
 242
 243         This dictionary corresponds to the GPT header format. Please, see the
 244         UEFI standard for the description of these fields.
 245
 246         If the 'primary' parameter is 'True', partitions from the primary GPT
 247         partition table are generated, otherwise partitions from the backup GPT
 248         partition table are generated. """
 249
 250         if primary:
 251             primary_str = "primary"
 252         else:
 253             primary_str = "backup"
 254
 255         header = self.read_header(primary)
 256         raw_ptable = self._read_raw_ptable(header)
 257
 258         for index in xrange(0, header['entries_cnt']):
 259             start = header['entry_size'] * index
 260             end = start + header['entry_size']
 261             raw_entry = struct.unpack(_GPT_ENTRY_FORMAT, raw_ptable[start:end])
 262
 263             if raw_entry[2] == 0 or raw_entry[3] == 0:
 264                 continue
 265
 266             part_name = str(raw_entry[5].decode('UTF-16').split('\0', 1)[0])
 267
 268             yield { 'index'       : index,
 269                     'offs'        : header['ptable_offs'] + start,
 270                     'type_uuid'   : _stringify_uuid(raw_entry[0]),
 271                     'part_uuid'   : _stringify_uuid(raw_entry[1]),
 272                     'first_lba'   : raw_entry[2],
 273                     'last_lba'    : raw_entry[3],
 274                     'flags'       : raw_entry[4],
 275                     'name'        : part_name,
 276                     'primary'     : primary,
 277                     'primary_str' : primary_str }
 278
 279     def _change_partition(self, header, entry):
 280         """ A helper function for 'change_partitions()' which changes a
 281         a paricular instance of the partition table (primary or backup). """
 282
 283         if entry['index'] >= header['entries_cnt']:
 284             raise MountError("Partition table at LBA %d has only %d "   \
 285                              "records cannot change record number " % \
 286                              (header['entries_cnt'], entry['index']))
 287         # Read raw GPT header
 288         raw_hdr = self._read_disk(header['hdr_offs'], _GPT_HEADER_SIZE)
 289         raw_hdr = list(struct.unpack(_GPT_HEADER_FORMAT, raw_hdr))
 290         _validate_header(raw_hdr)
 291
 292         # Prepare the new partition table entry
 293         raw_entry = struct.pack(_GPT_ENTRY_FORMAT,
 294                                 uuid.UUID(entry['type_uuid']).bytes_le,
 295                                 uuid.UUID(entry['part_uuid']).bytes_le,
 296                                 entry['first_lba'],
 297                                 entry['last_lba'],
 298                                 entry['flags'],
 299                                 entry['name'].encode('UTF-16LE'))
 300
 301         # Write the updated entry to the disk
 302         entry_offs = header['ptable_offs'] + \
 303                      header['entry_size'] * entry['index']
 304         self._write_disk(entry_offs, raw_entry)
 305
 306         # Calculate and update partition table CRC32
 307         raw_ptable = self._read_disk(header['ptable_offs'],
 308                                      header['ptable_size'])
 309         raw_hdr[13] = binascii.crc32(raw_ptable) & 0xFFFFFFFF
 310
 311         # Calculate and update the GPT header CRC
 312         raw_hdr[3] = _calc_header_crc(raw_hdr)
 313
 314         # Write the updated header to the disk
 315         raw_hdr = struct.pack(_GPT_HEADER_FORMAT, *raw_hdr)
 316         self._write_disk(header['hdr_offs'], raw_hdr)
 317
 318     def change_partition(self, entry):
 319         """ Change a GPT partition. The 'entry' argument has the same format as
 320         'get_partitions()' returns. This function simply changes the partition
 321         table record corresponding to 'entry' in both, the primary and the
 322         backup GPT partition tables. The parition table CRC is re-calculated
 323         and the GPT headers are modified accordingly. """
 324
 325         # Change the primary partition table
 326         header = self.read_header(True)
 327         self._change_partition(header, entry)
 328
 329         # Change the backup partition table
 330         header = self.read_header(False)
 331         self._change_partition(header, entry)