code.grnet.gr Git - snf-image/blob - snf-image-helper/fix_gpt.py

   1 #!/usr/bin/env python
   2
   3 # Copyright (C) 2012 GRNET S.A.
   4 #
   5 # This program is free software; you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation; either version 2 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful, but
  11 # WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  13 # General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with this program; if not, write to the Free Software
  17 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  18 # 02110-1301, USA.
  19
  20 import struct
  21 import sys
  22 import binascii
  23 import stat
  24 import os
  25
  26 BLOCKSIZE = 512
  27
  28
  29 class MBR(object):
  30     """Represents a Master Boot Record."""
  31     class Partition(object):
  32         format = "<B3sB3sLL"
  33
  34         def __init__(self, raw_part):
  35             (
  36                 self.status,
  37                 self.start,
  38                 self.type,
  39                 self.end,
  40                 self.first_sector,
  41                 self.sector_count
  42             ) = struct.unpack(self.format, raw_part)
  43
  44         def pack(self):
  45             return struct.pack(self.format,
  46                 self.status,
  47                 self.start,
  48                 self.type,
  49                 self.end,
  50                 self.first_sector,
  51                 self.sector_count
  52             )
  53
  54         @staticmethod
  55         def size():
  56             """Returns the size of an MBR partition entry"""
  57             return struct.calcsize(MBR.Partition.format)
  58
  59         def unpack_chs(self, chs):
  60             """Unpacks a CHS address string to a tuple."""
  61
  62             assert len(chs) == 3
  63
  64             head = struct.unpack('<B', chs[0])[0]
  65             sector = struct.unpack('<B', chs[1])[0] & 0x3f
  66             cylinder = (struct.unpack('<B', chs[1])[0] & 0xC0) << 2 | \
  67                 struct.unpack('<B', chs[2])[0]
  68
  69             return (cylinder, head, sector)
  70
  71         def pack_chs(self, cylinder, head, sector):
  72             """Packs a CHS tuple to an address string."""
  73
  74             assert 1 <= sector <= 63
  75             assert 0 <= cylinder <= 1023
  76             assert 0 <= head <= 255
  77
  78             byte0 = head
  79             byte1 = (cylinder >> 2) & 0xC0 | sector
  80             byte2 = cylinder & 0xff
  81
  82             return struct.pack('<BBB', byte0, byte1, byte2)
  83
  84     format = "<444s2x16s16s16s16s2s"
  85     """
  86     Offset  Length          Contents
  87     0       440(max. 446)   code area
  88     440     2(optional)     disk signature
  89     444     2               Usually nulls
  90     446     16              Partition 0
  91     462     16              Partition 1
  92     478     16              Partition 2
  93     494     16              Partition 3
  94     510     2               MBR signature
  95     """
  96     def __init__(self, block):
  97         raw_part = {}
  98         self.code_area, \
  99         raw_part[0], \
 100         raw_part[1], \
 101         raw_part[2], \
 102         raw_part[3], \
 103         self.signature = struct.unpack(self.format, block)
 104
 105         self.part = {}
 106         for i in range(4):
 107             self.part[i] = self.Partition(raw_part[i])
 108
 109     @staticmethod
 110     def size():
 111         """Returns the size of a Master Boot Record."""
 112         return struct.calcsize(MBR.format)
 113
 114     def pack(self):
 115         """Packs an MBR to a binary string."""
 116         return struct.pack(self.format,
 117             self.code_area,
 118             self.part[0].pack(),
 119             self.part[1].pack(),
 120             self.part[2].pack(),
 121             self.part[3].pack(),
 122             self.signature
 123         )
 124
 125
 126 class GPTPartitionTable(object):
 127     """Represents a GUID Partition Table."""
 128     class GPTHeader(object):
 129         """Represents a GPT Header of a GUID Partition Table."""
 130         format = "<8s4sII4xQQQQ16sQIII"
 131         """
 132         Offset  Length          Contents
 133         0       8 bytes         Signature
 134         8       4 bytes         Revision
 135         12      4 bytes         Header size in little endian
 136         16      4 bytes         CRC32 of header
 137         20      4 bytes         Reserved; must be zero
 138         24      8 bytes         Current LBA
 139         32      8 bytes         Backup LBA
 140         40      8 bytes         First usable LBA for partitions
 141         48      8 bytes         Last usable LBA
 142         56      16 bytes        Disk GUID
 143         72      8 bytes         Partition entries starting LBA
 144         80      4 bytes         Number of partition entries
 145         84      4 bytes         Size of a partition entry
 146         88      4 bytes         CRC32 of partition array
 147         92      *               Reserved; must be zeroes
 148         LBA    size            Total
 149         """
 150
 151         def __init__(self, block):
 152             self.signature, \
 153             self.revision, \
 154             self.hdr_size, \
 155             self.header_crc32, \
 156             self.current_lba, \
 157             self.backup_lba, \
 158             self.first_usable_lba, \
 159             self.last_usable_lba, \
 160             self.uuid, \
 161             self.part_entry_start, \
 162             self.part_count, \
 163             self.part_entry_size, \
 164             self.part_crc32 = struct.unpack(self.format, block)
 165
 166         def pack(self):
 167             """Packs a GPT Header to a binary string."""
 168             return struct.pack(self.format,
 169                 self.signature, \
 170                 self.revision, \
 171                 self.hdr_size, \
 172                 self.header_crc32, \
 173                 self.current_lba, \
 174                 self.backup_lba, \
 175                 self.first_usable_lba, \
 176                 self.last_usable_lba, \
 177                 self.uuid, \
 178                 self.part_entry_start, \
 179                 self.part_count, \
 180                 self.part_entry_size, \
 181                 self.part_crc32
 182             )
 183
 184         @staticmethod
 185         def size():
 186             """Returns the size of a GPT Header."""
 187             return struct.calcsize(GPTPartitionTable.GPTHeader.format)
 188
 189     def __init__(self, disk):
 190         self.disk = disk
 191         with open(disk, "rb") as d:
 192             # MBR (Logical block address 0)
 193             lba0 = d.read(BLOCKSIZE)
 194             self.mbr = MBR(lba0)
 195
 196             # Primary GPT Header (LBA 1)
 197             raw_header = d.read(self.GPTHeader.size())
 198             self.primary = self.GPTHeader(raw_header)
 199
 200             # Partition entries (LBA 2...34)
 201             d.seek(self.primary.part_entry_start * BLOCKSIZE)
 202             entries_size = self.primary.part_count * \
 203                                                 self.primary.part_entry_size
 204             self.part_entries = d.read(entries_size)
 205
 206             # Secondary GPT Header (LBA -1)
 207             d.seek(self.primary.backup_lba * BLOCKSIZE)
 208             raw_header = d.read(self.GPTHeader.size())
 209             self.secondary = self.GPTHeader(raw_header)
 210
 211     def size(self):
 212         """Returns the payload size of GPT partitioned device."""
 213         return (self.primary.backup_lba + 1) * BLOCKSIZE
 214
 215     def fix(self, lba_count):
 216         """Move the secondary GPT Header entries to the LBA specified by
 217         lba_count parameter.
 218         """
 219
 220         assert lba_count * BLOCKSIZE > self.size()
 221
 222         # Correct MBR
 223         #TODO: Check if the partition tables is hybrid
 224         self.mbr.part[0].sector_count = lba_count - 1
 225
 226         # Fix Primary header
 227         self.primary.header_crc32 = 0
 228         self.primary.backup_lba = lba_count - 1  # LBA-1
 229         self.primary.last_usable_lba = lba_count - 34  # LBA-34
 230         self.primary.header_crc32 = \
 231                             binascii.crc32(self.primary.pack()) & 0xffffffff
 232
 233         # Fix Secondary header
 234         self.secondary.header_crc32 = 0
 235         self.secondary.current_lba = self.primary.backup_lba
 236         self.secondary.last_usable_lba = lba_count - 34  # LBA-34
 237         self.secondary.part_entry_start = lba_count - 33  # LBA-33
 238         self.secondary.header_crc32 = \
 239                             binascii.crc32(self.secondary.pack()) & 0xffffffff
 240
 241         # Copy the new partition table back to the device
 242         with open(self.disk, "wb") as d:
 243             d.write(self.mbr.pack())
 244             d.write(self.primary.pack())
 245             d.write('\x00' * (BLOCKSIZE - self.primary.size()))
 246             d.seek(self.secondary.part_entry_start * BLOCKSIZE)
 247             d.write(self.part_entries)
 248             d.seek(self.primary.backup_lba * BLOCKSIZE)
 249             d.write(self.secondary.pack())
 250             d.write('\x00' * (BLOCKSIZE - self.secondary.size()))
 251
 252
 253 if __name__ == '__main__':
 254     usage = "Usage: %s <disk> <sectors>\n" % (sys.argv[0])
 255
 256     if len(sys.argv) != 3:
 257         sys.stderr.write(usage)
 258         sys.exit(1)
 259
 260     disk = sys.argv[1]
 261     mode = os.stat(disk).st_mode
 262     if not stat.S_ISBLK(mode):
 263         sys.stderr.write("Parameter disk must be a block device\n")
 264         sys.stderr.write(usage)
 265         sys.exit(1)
 266
 267     try:
 268         size = int(sys.argv[2])
 269     except ValueError:
 270         sys.stderr.write("Parameter new_size must be a number\n")
 271         sys.stderr.write(usage)
 272         sys.exit(1)
 273
 274     ptable = GPTPartitionTable(disk)
 275     if size * BLOCKSIZE == ptable.size():
 276         sys.stderr.write("Nothing to do...\n")
 277     elif size * BLOCKSIZE > ptable.size():
 278         ptable.fix(size)
 279         sys.stderr.write("GPT table was fixed\n")
 280     else:
 281         sys.stderr.write("Disk is langer than size")
 282         exit(1)
 283
 284     sys.exit(0)
 285
 286 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :