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

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