You-Cheng Syu | d569294 | 2018-01-04 14:40:59 +0800 | [diff] [blame] | 1 | #!/usr/bin/env python |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 2 | # Copyright 2017 The Chromium OS Authors. All rights reserved. |
| 3 | # Use of this source code is governed by a BSD-style license that can be |
| 4 | # found in the LICENSE file. |
| 5 | |
| 6 | """An utility to manipulate GPT on a disk image. |
| 7 | |
| 8 | Chromium OS factory software usually needs to access partitions from disk |
| 9 | images. However, there is no good, lightweight, and portable GPT utility. |
| 10 | Most Chromium OS systems use `cgpt`, but that's not by default installed on |
| 11 | Ubuntu. Most systems have parted (GNU) or partx (util-linux-ng) but they have |
| 12 | their own problems. |
| 13 | |
| 14 | For example, when a disk image is resized (usually enlarged for putting more |
| 15 | resources on stateful partition), GPT table must be updated. However, |
| 16 | - `parted` can't repair partition without interactive console in exception |
| 17 | handler. |
| 18 | - `partx` cannot fix headers nor make changes to partition table. |
| 19 | - `cgpt repair` does not fix `LastUsableLBA` so we cannot enlarge partition. |
| 20 | - `gdisk` is not installed on most systems. |
| 21 | |
| 22 | As a result, we need a dedicated tool to help processing GPT. |
| 23 | |
| 24 | This pygpt.py provides a simple and customized implementation for processing |
| 25 | GPT, as a replacement for `cgpt`. |
| 26 | """ |
| 27 | |
| 28 | |
| 29 | from __future__ import print_function |
| 30 | |
| 31 | import argparse |
| 32 | import binascii |
| 33 | import collections |
| 34 | import logging |
| 35 | import os |
| 36 | import struct |
| 37 | import uuid |
| 38 | |
| 39 | |
| 40 | # The binascii.crc32 returns signed integer, so CRC32 in in struct must be |
| 41 | # declared as 'signed' (l) instead of 'unsigned' (L). |
| 42 | # http://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_table_header_.28LBA_1.29 |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 43 | HEADER_DESCRIPTION = """ |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 44 | 8s Signature |
| 45 | 4s Revision |
| 46 | L HeaderSize |
| 47 | l CRC32 |
| 48 | 4s Reserved |
| 49 | Q CurrentLBA |
| 50 | Q BackupLBA |
| 51 | Q FirstUsableLBA |
| 52 | Q LastUsableLBA |
| 53 | 16s DiskGUID |
| 54 | Q PartitionEntriesStartingLBA |
| 55 | L PartitionEntriesNumber |
| 56 | L PartitionEntrySize |
| 57 | l PartitionArrayCRC32 |
| 58 | """ |
| 59 | |
| 60 | # http://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 61 | PARTITION_DESCRIPTION = """ |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 62 | 16s TypeGUID |
| 63 | 16s UniqueGUID |
| 64 | Q FirstLBA |
| 65 | Q LastLBA |
| 66 | Q Attributes |
| 67 | 72s Names |
| 68 | """ |
| 69 | |
| 70 | |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 71 | def BitProperty(getter, setter, shift, mask): |
| 72 | """A generator for bit-field properties. |
| 73 | |
| 74 | This is used inside a class to manipulate an integer-like variable using |
| 75 | properties. The getter and setter should be member functions to change the |
| 76 | underlying member data. |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 77 | |
| 78 | Args: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 79 | getter: a function to read integer type variable (for all the bits). |
| 80 | setter: a function to set the new changed integer type variable. |
| 81 | shift: integer for how many bits should be shifted (right). |
| 82 | mask: integer for the mask to filter out bit field. |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 83 | """ |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 84 | def _getter(self): |
| 85 | return (getter(self) >> shift) & mask |
| 86 | def _setter(self, value): |
| 87 | assert value & mask == value, ( |
| 88 | 'Value %s out of range (mask=%s)' % (value, mask)) |
| 89 | setter(self, getter(self) & ~(mask << shift) | value << shift) |
| 90 | return property(_getter, _setter) |
| 91 | |
| 92 | |
| 93 | class GPTBlob(object): |
| 94 | """A decorator class to help accessing GPT blobs as named tuple. |
| 95 | |
| 96 | To use this, specify the blob description (struct format and named tuple field |
| 97 | names) above the derived class, for example: |
| 98 | |
| 99 | @GPTBlob(description): |
| 100 | class Header(GPTObject): |
| 101 | pass |
| 102 | """ |
| 103 | def __init__(self, description): |
| 104 | spec = description.split() |
| 105 | self.struct_format = '<' + ''.join(spec[::2]) |
| 106 | self.fields = spec[1::2] |
| 107 | |
| 108 | def __call__(self, cls): |
| 109 | new_bases = (( |
| 110 | collections.namedtuple(cls.__name__, self.fields),) + cls.__bases__) |
| 111 | new_cls = type(cls.__name__, new_bases, dict(cls.__dict__)) |
| 112 | setattr(new_cls, 'FORMAT', self.struct_format) |
| 113 | return new_cls |
| 114 | |
| 115 | |
| 116 | class GPTObject(object): |
| 117 | """An object in GUID Partition Table. |
| 118 | |
| 119 | This needs to be decorated by @GPTBlob(description) and inherited by a real |
| 120 | class. Properties (not member functions) in CamelCase should be reserved for |
| 121 | named tuple attributes. |
| 122 | |
| 123 | To create a new object, use class method ReadFrom(), which takes a stream |
| 124 | as input or None to create with all elements set to zero. To make changes to |
| 125 | named tuple elements, use member function Clone(changes). |
| 126 | |
| 127 | It is also possible to attach some additional properties to the object as meta |
| 128 | data (for example path of the underlying image file). To do that, specify the |
| 129 | data as keyword arguments when calling ReadFrom(). These properties will be |
| 130 | preserved when you call Clone(). |
| 131 | |
| 132 | A special case is "reset named tuple elements of an object but keeping all |
| 133 | properties", for example changing a partition object to unused (zeroed). |
| 134 | ReadFrom() is a class method so properties won't be copied. You need to |
| 135 | call as cls.ReadFrom(None, **p.__dict__), or a short cut - p.CloneAndZero(). |
| 136 | """ |
| 137 | |
| 138 | FORMAT = None |
| 139 | """The struct.{pack,unpack} format string, and should be set by GPTBlob.""" |
| 140 | |
| 141 | CLONE_CONVERTERS = None |
| 142 | """A dict (name, cvt) to convert input arguments into named tuple data. |
| 143 | |
| 144 | `name` is a string for the name of argument to convert. |
| 145 | `cvt` is a callable to convert value. The return value may be: |
| 146 | - a tuple in (new_name, value): save the value as new name. |
| 147 | - otherwise, save the value in original name. |
| 148 | Note tuple is an invalid input for struct.unpack so it's used for the |
| 149 | special value. |
| 150 | """ |
| 151 | |
| 152 | @classmethod |
| 153 | def ReadFrom(cls, f, **kargs): |
| 154 | """Reads and decode an object from stream. |
| 155 | |
| 156 | Args: |
| 157 | f: a stream to read blob, or None to decode with all zero bytes. |
| 158 | kargs: a dict for additional attributes in object. |
| 159 | """ |
| 160 | if f is None: |
| 161 | reader = lambda num: '\x00' * num |
| 162 | else: |
| 163 | reader = f.read |
| 164 | data = cls(*struct.unpack(cls.FORMAT, reader(struct.calcsize(cls.FORMAT)))) |
| 165 | # Named tuples do not accept kargs in constructor. |
| 166 | data.__dict__.update(kargs) |
| 167 | return data |
| 168 | |
| 169 | def Clone(self, **dargs): |
| 170 | """Clones a new instance with modifications. |
| 171 | |
| 172 | GPT objects are usually named tuples that are immutable, so the only way |
| 173 | to make changes is to create a new instance with modifications. |
| 174 | |
| 175 | Args: |
| 176 | dargs: a dict with all modifications. |
| 177 | """ |
| 178 | for name, convert in (self.CLONE_CONVERTERS or {}).iteritems(): |
| 179 | if name not in dargs: |
| 180 | continue |
| 181 | result = convert(dargs.pop(name)) |
| 182 | if isinstance(result, tuple): |
| 183 | assert len(result) == 2, 'Converted tuple must be (name, value).' |
| 184 | dargs[result[0]] = result[1] |
| 185 | else: |
| 186 | dargs[name] = result |
| 187 | |
| 188 | cloned = self._replace(**dargs) |
| 189 | cloned.__dict__.update(self.__dict__) |
| 190 | return cloned |
| 191 | |
| 192 | def CloneAndZero(self, **dargs): |
| 193 | """Short cut to create a zeroed object while keeping all properties. |
| 194 | |
| 195 | This is very similar to Clone except all named tuple elements will be zero. |
| 196 | Also different from class method ReadFrom(None) because this keeps all |
| 197 | properties from one object. |
| 198 | """ |
| 199 | cloned = self.ReadFrom(None, **self.__dict__) |
| 200 | return cloned.Clone(**dargs) if dargs else cloned |
| 201 | |
| 202 | @property |
| 203 | def blob(self): |
| 204 | """Returns the object in formatted bytes.""" |
| 205 | return struct.pack(self.FORMAT, *self) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 206 | |
| 207 | |
| 208 | class GPT(object): |
| 209 | """A GPT helper class. |
| 210 | |
| 211 | To load GPT from an existing disk image file, use `LoadFromFile`. |
| 212 | After modifications were made, use `WriteToFile` to commit changes. |
| 213 | |
| 214 | Attributes: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 215 | header: a namedtuple of GPT header. |
| 216 | partitions: a list of GPT partition entry nametuple. |
| 217 | block_size: integer for size of bytes in one block (sector). |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 218 | """ |
| 219 | |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 220 | DEFAULT_BLOCK_SIZE = 512 |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 221 | TYPE_GUID_UNUSED = '\x00' * 16 |
| 222 | TYPE_GUID_MAP = { |
| 223 | '00000000-0000-0000-0000-000000000000': 'Unused', |
| 224 | 'EBD0A0A2-B9E5-4433-87C0-68B6B72699C7': 'Linux data', |
| 225 | 'FE3A2A5D-4F32-41A7-B725-ACCC3285A309': 'ChromeOS kernel', |
| 226 | '3CB8E202-3B7E-47DD-8A3C-7FF2A13CFCEC': 'ChromeOS rootfs', |
| 227 | '2E0A753D-9E48-43B0-8337-B15192CB1B5E': 'ChromeOS reserved', |
| 228 | 'CAB6E88E-ABF3-4102-A07A-D4BB9BE3C1D3': 'ChromeOS firmware', |
| 229 | 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B': 'EFI System Partition', |
| 230 | } |
| 231 | TYPE_GUID_LIST_BOOTABLE = [ |
| 232 | 'FE3A2A5D-4F32-41A7-B725-ACCC3285A309', # ChromeOS kernel |
| 233 | 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B', # EFI System Partition |
| 234 | ] |
| 235 | |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 236 | @GPTBlob(HEADER_DESCRIPTION) |
| 237 | class Header(GPTObject): |
| 238 | """Wrapper to Header in GPT.""" |
| 239 | SIGNATURES = ['EFI PART', 'CHROMEOS'] |
| 240 | SIGNATURE_IGNORE = 'IGNOREME' |
| 241 | DEFAULT_REVISION = '\x00\x00\x01\x00' |
| 242 | |
| 243 | DEFAULT_PARTITION_ENTRIES = 128 |
| 244 | DEFAULT_PARTITIONS_LBA = 2 # LBA 0 = MBR, LBA 1 = GPT Header. |
| 245 | |
| 246 | def Clone(self, **dargs): |
| 247 | """Creates a new instance with modifications. |
| 248 | |
| 249 | GPT objects are usually named tuples that are immutable, so the only way |
| 250 | to make changes is to create a new instance with modifications. |
| 251 | |
| 252 | CRC32 is always updated but PartitionArrayCRC32 must be updated explicitly |
| 253 | since we can't track changes in GPT.partitions automatically. |
| 254 | |
| 255 | Note since GPTHeader.Clone will always update CRC, we can only check and |
| 256 | compute CRC by super(GPT.Header, header).Clone, or header._replace. |
| 257 | """ |
| 258 | dargs['CRC32'] = 0 |
| 259 | header = super(GPT.Header, self).Clone(**dargs) |
| 260 | return super(GPT.Header, header).Clone(CRC32=binascii.crc32(header.blob)) |
| 261 | |
| 262 | class PartitionAttributes(object): |
| 263 | """Wrapper for Partition.Attributes. |
| 264 | |
| 265 | This can be created using Partition.attrs, but the changed properties won't |
| 266 | apply to underlying Partition until an explicit call with |
| 267 | Partition.Clone(Attributes=new_attrs). |
| 268 | """ |
| 269 | |
| 270 | def __init__(self, attrs): |
| 271 | self._attrs = attrs |
| 272 | |
| 273 | @property |
| 274 | def raw(self): |
| 275 | """Returns the raw integer type attributes.""" |
| 276 | return self._Get() |
| 277 | |
| 278 | def _Get(self): |
| 279 | return self._attrs |
| 280 | |
| 281 | def _Set(self, value): |
| 282 | self._attrs = value |
| 283 | |
| 284 | successful = BitProperty(_Get, _Set, 56, 1) |
| 285 | tries = BitProperty(_Get, _Set, 52, 0xf) |
| 286 | priority = BitProperty(_Get, _Set, 48, 0xf) |
| 287 | legacy_boot = BitProperty(_Get, _Set, 2, 1) |
| 288 | required = BitProperty(_Get, _Set, 0, 1) |
| 289 | |
| 290 | @GPTBlob(PARTITION_DESCRIPTION) |
| 291 | class Partition(GPTObject): |
| 292 | """The partition entry in GPT. |
| 293 | |
| 294 | Please include following properties when creating a Partition object: |
| 295 | - image: a string for path to the image file the partition maps to. |
| 296 | - number: the 1-based partition number. |
| 297 | - block_size: an integer for size of each block (LBA, or sector). |
| 298 | """ |
| 299 | NAMES_ENCODING = 'utf-16-le' |
| 300 | NAMES_LENGTH = 72 |
| 301 | |
| 302 | CLONE_CONVERTERS = { |
| 303 | # TODO(hungte) check if encoded name is too long. |
| 304 | 'label': lambda l: (None if l is None else |
| 305 | ('Names', l.encode(GPT.Partition.NAMES_ENCODING))), |
| 306 | 'TypeGUID': lambda v: v.bytes_le if isinstance(v, uuid.UUID) else v, |
| 307 | 'UniqueGUID': lambda v: v.bytes_le if isinstance(v, uuid.UUID) else v, |
| 308 | 'Attributes': ( |
| 309 | lambda v: v.raw if isinstance(v, GPT.PartitionAttributes) else v), |
| 310 | } |
| 311 | |
| 312 | def __str__(self): |
| 313 | return '%s#%s' % (self.image, self.number) |
| 314 | |
| 315 | def IsUnused(self): |
| 316 | """Returns if the partition is unused and can be allocated.""" |
| 317 | return self.TypeGUID == GPT.TYPE_GUID_UNUSED |
| 318 | |
| 319 | @property |
| 320 | def blocks(self): |
| 321 | """Return size of partition in blocks (see block_size).""" |
| 322 | return self.LastLBA - self.FirstLBA + 1 |
| 323 | |
| 324 | @property |
| 325 | def offset(self): |
| 326 | """Returns offset to partition in bytes.""" |
| 327 | return self.FirstLBA * self.block_size |
| 328 | |
| 329 | @property |
| 330 | def size(self): |
| 331 | """Returns size of partition in bytes.""" |
| 332 | return self.blocks * self.block_size |
| 333 | |
| 334 | @property |
| 335 | def type_guid(self): |
| 336 | return uuid.UUID(bytes_le=self.TypeGUID) |
| 337 | |
| 338 | @property |
| 339 | def unique_guid(self): |
| 340 | return uuid.UUID(bytes_le=self.UniqueGUID) |
| 341 | |
| 342 | @property |
| 343 | def label(self): |
| 344 | """Returns the Names in decoded string type.""" |
| 345 | return self.Names.decode(self.NAMES_ENCODING).strip('\0') |
| 346 | |
| 347 | @property |
| 348 | def attrs(self): |
| 349 | return GPT.PartitionAttributes(self.Attributes) |
| 350 | |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 351 | def __init__(self): |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 352 | """GPT constructor. |
| 353 | |
| 354 | See LoadFromFile for how it's usually used. |
| 355 | """ |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 356 | self.header = None |
| 357 | self.partitions = None |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 358 | self.block_size = self.DEFAULT_BLOCK_SIZE |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 359 | |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 360 | @classmethod |
Hung-Te Lin | 6977ae1 | 2018-04-17 12:20:32 +0800 | [diff] [blame] | 361 | def LoadFromFile(cls, image): |
| 362 | """Loads a GPT table from give disk image file object. |
| 363 | |
| 364 | Args: |
| 365 | image: a string as file path or a file-like object to read from. |
| 366 | """ |
| 367 | if isinstance(image, basestring): |
| 368 | with open(image, 'rb') as f: |
| 369 | return cls.LoadFromFile(f) |
| 370 | |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 371 | gpt = cls() |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 372 | # Try DEFAULT_BLOCK_SIZE, then 4K. |
| 373 | for block_size in [cls.DEFAULT_BLOCK_SIZE, 4096]: |
Hung-Te Lin | 6977ae1 | 2018-04-17 12:20:32 +0800 | [diff] [blame] | 374 | image.seek(block_size * 1) |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 375 | header = gpt.Header.ReadFrom(image) |
| 376 | if header.Signature in cls.Header.SIGNATURES: |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 377 | gpt.block_size = block_size |
| 378 | break |
| 379 | else: |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 380 | raise ValueError('Invalid signature in GPT header.') |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 381 | |
Hung-Te Lin | 6977ae1 | 2018-04-17 12:20:32 +0800 | [diff] [blame] | 382 | image.seek(gpt.block_size * header.PartitionEntriesStartingLBA) |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 383 | def ReadPartition(image, i): |
| 384 | p = gpt.Partition.ReadFrom( |
| 385 | image, image=image.name, number=i + 1, block_size=gpt.block_size) |
| 386 | return p |
| 387 | |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 388 | gpt.header = header |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 389 | gpt.partitions = [ |
| 390 | ReadPartition(image, i) for i in range(header.PartitionEntriesNumber)] |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 391 | return gpt |
| 392 | |
| 393 | def GetValidPartitions(self): |
| 394 | """Returns the list of partitions before entry with empty type GUID. |
| 395 | |
| 396 | In partition table, the first entry with empty type GUID indicates end of |
| 397 | valid partitions. In most implementations all partitions after that should |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 398 | be zeroed. However, few implementations for example cgpt, may create |
| 399 | partitions in arbitrary order so use this carefully. |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 400 | """ |
| 401 | for i, p in enumerate(self.partitions): |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 402 | if p.IsUnused(): |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 403 | return self.partitions[:i] |
| 404 | return self.partitions |
| 405 | |
| 406 | def GetMaxUsedLBA(self): |
Hung-Te Lin | d3a2e9a | 2018-04-19 13:07:26 +0800 | [diff] [blame] | 407 | """Returns the max LastLBA from all used partitions.""" |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 408 | parts = [p for p in self.partitions if not p.IsUnused()] |
Hung-Te Lin | d3a2e9a | 2018-04-19 13:07:26 +0800 | [diff] [blame] | 409 | return (max(p.LastLBA for p in parts) |
| 410 | if parts else self.header.FirstUsableLBA - 1) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 411 | |
| 412 | def GetPartitionTableBlocks(self, header=None): |
| 413 | """Returns the blocks (or LBA) of partition table from given header.""" |
| 414 | if header is None: |
| 415 | header = self.header |
| 416 | size = header.PartitionEntrySize * header.PartitionEntriesNumber |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 417 | blocks = size / self.block_size |
| 418 | if size % self.block_size: |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 419 | blocks += 1 |
| 420 | return blocks |
| 421 | |
| 422 | def Resize(self, new_size): |
| 423 | """Adjust GPT for a disk image in given size. |
| 424 | |
| 425 | Args: |
| 426 | new_size: Integer for new size of disk image file. |
| 427 | """ |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 428 | old_size = self.block_size * (self.header.BackupLBA + 1) |
| 429 | if new_size % self.block_size: |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 430 | raise ValueError('New file size %d is not valid for image files.' % |
| 431 | new_size) |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 432 | new_blocks = new_size / self.block_size |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 433 | if old_size != new_size: |
| 434 | logging.warn('Image size (%d, LBA=%d) changed from %d (LBA=%d).', |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 435 | new_size, new_blocks, old_size, old_size / self.block_size) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 436 | else: |
| 437 | logging.info('Image size (%d, LBA=%d) not changed.', |
| 438 | new_size, new_blocks) |
Hung-Te Lin | d3a2e9a | 2018-04-19 13:07:26 +0800 | [diff] [blame] | 439 | return |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 440 | |
Hung-Te Lin | d3a2e9a | 2018-04-19 13:07:26 +0800 | [diff] [blame] | 441 | # Expected location |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 442 | backup_lba = new_blocks - 1 |
Hung-Te Lin | d3a2e9a | 2018-04-19 13:07:26 +0800 | [diff] [blame] | 443 | last_usable_lba = backup_lba - self.header.FirstUsableLBA |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 444 | |
Hung-Te Lin | d3a2e9a | 2018-04-19 13:07:26 +0800 | [diff] [blame] | 445 | if last_usable_lba < self.header.LastUsableLBA: |
| 446 | max_used_lba = self.GetMaxUsedLBA() |
| 447 | if last_usable_lba < max_used_lba: |
| 448 | raise ValueError('Backup partition tables will overlap used partitions') |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 449 | |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 450 | self.header = self.header.Clone( |
| 451 | BackupLBA=backup_lba, LastUsableLBA=last_usable_lba) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 452 | |
| 453 | def GetFreeSpace(self): |
| 454 | """Returns the free (available) space left according to LastUsableLBA.""" |
| 455 | max_lba = self.GetMaxUsedLBA() |
| 456 | assert max_lba <= self.header.LastUsableLBA, "Partitions too large." |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 457 | return self.block_size * (self.header.LastUsableLBA - max_lba) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 458 | |
| 459 | def ExpandPartition(self, i): |
| 460 | """Expands a given partition to last usable LBA. |
| 461 | |
| 462 | Args: |
| 463 | i: Index (0-based) of target partition. |
| 464 | """ |
| 465 | # Assume no partitions overlap, we need to make sure partition[i] has |
| 466 | # largest LBA. |
| 467 | if i < 0 or i >= len(self.GetValidPartitions()): |
| 468 | raise ValueError('Partition index %d is invalid.' % (i + 1)) |
| 469 | p = self.partitions[i] |
| 470 | max_used_lba = self.GetMaxUsedLBA() |
| 471 | if max_used_lba > p.LastLBA: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 472 | raise ValueError( |
| 473 | 'Cannot expand partition %d because it is not the last allocated ' |
| 474 | 'partition.' % (i + 1)) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 475 | |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 476 | old_blocks = p.blocks |
| 477 | p = p.Clone(LastLBA=self.header.LastUsableLBA) |
| 478 | new_blocks = p.blocks |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 479 | self.partitions[i] = p |
| 480 | logging.warn('Partition NR=%d expanded, size in LBA: %d -> %d.', |
| 481 | i + 1, old_blocks, new_blocks) |
| 482 | |
| 483 | def UpdateChecksum(self): |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 484 | """Updates all checksum fields in GPT objects. |
| 485 | |
| 486 | The Header.CRC32 is automatically updated in Header.Clone(). |
| 487 | """ |
| 488 | parts = ''.join(p.blob for p in self.partitions) |
| 489 | self.header = self.header.Clone( |
| 490 | PartitionArrayCRC32=binascii.crc32(parts)) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 491 | |
| 492 | def GetBackupHeader(self): |
| 493 | """Returns the backup header according to current header.""" |
| 494 | partitions_starting_lba = ( |
| 495 | self.header.BackupLBA - self.GetPartitionTableBlocks()) |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 496 | return self.header.Clone( |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 497 | BackupLBA=self.header.CurrentLBA, |
| 498 | CurrentLBA=self.header.BackupLBA, |
| 499 | PartitionEntriesStartingLBA=partitions_starting_lba) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 500 | |
Hung-Te Lin | 6977ae1 | 2018-04-17 12:20:32 +0800 | [diff] [blame] | 501 | def WriteToFile(self, image): |
| 502 | """Updates partition table in a disk image file. |
| 503 | |
| 504 | Args: |
| 505 | image: a string as file path or a file-like object to write into. |
| 506 | """ |
| 507 | if isinstance(image, basestring): |
| 508 | with open(image, 'rb+') as f: |
| 509 | return self.WriteToFile(f) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 510 | |
| 511 | def WriteData(name, blob, lba): |
| 512 | """Writes a blob into given location.""" |
| 513 | logging.info('Writing %s in LBA %d (offset %d)', |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 514 | name, lba, lba * self.block_size) |
Hung-Te Lin | 6977ae1 | 2018-04-17 12:20:32 +0800 | [diff] [blame] | 515 | image.seek(lba * self.block_size) |
| 516 | image.write(blob) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 517 | |
| 518 | self.UpdateChecksum() |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 519 | parts_blob = ''.join(p.blob for p in self.partitions) |
| 520 | WriteData('GPT Header', self.header.blob, self.header.CurrentLBA) |
| 521 | WriteData( |
| 522 | 'GPT Partitions', parts_blob, self.header.PartitionEntriesStartingLBA) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 523 | logging.info('Usable LBA: First=%d, Last=%d', |
| 524 | self.header.FirstUsableLBA, self.header.LastUsableLBA) |
| 525 | backup_header = self.GetBackupHeader() |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 526 | WriteData( |
| 527 | 'Backup Partitions', parts_blob, |
| 528 | backup_header.PartitionEntriesStartingLBA) |
| 529 | WriteData('Backup Header', backup_header.blob, backup_header.CurrentLBA) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 530 | |
| 531 | |
| 532 | class GPTCommands(object): |
| 533 | """Collection of GPT sub commands for command line to use. |
| 534 | |
| 535 | The commands are derived from `cgpt`, but not necessary to be 100% compatible |
| 536 | with cgpt. |
| 537 | """ |
| 538 | |
| 539 | FORMAT_ARGS = [ |
Peter Shih | c7156ca | 2018-02-26 14:46:24 +0800 | [diff] [blame] | 540 | ('begin', 'beginning sector'), |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 541 | ('size', 'partition size (in sectors)'), |
Peter Shih | c7156ca | 2018-02-26 14:46:24 +0800 | [diff] [blame] | 542 | ('type', 'type guid'), |
| 543 | ('unique', 'unique guid'), |
| 544 | ('label', 'label'), |
| 545 | ('Successful', 'Successful flag'), |
| 546 | ('Tries', 'Tries flag'), |
| 547 | ('Priority', 'Priority flag'), |
| 548 | ('Legacy', 'Legacy Boot flag'), |
| 549 | ('Attribute', 'raw 16-bit attribute value (bits 48-63)')] |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 550 | |
| 551 | def __init__(self): |
| 552 | pass |
| 553 | |
| 554 | @classmethod |
| 555 | def RegisterRepair(cls, p): |
| 556 | """Registers the repair command to argparser. |
| 557 | |
| 558 | Args: |
| 559 | p: An argparse parser instance. |
| 560 | """ |
| 561 | p.add_argument( |
| 562 | '--expand', action='store_true', default=False, |
| 563 | help='Expands stateful partition to full disk.') |
| 564 | p.add_argument('image_file', type=argparse.FileType('rb+'), |
| 565 | help='Disk image file to repair.') |
| 566 | |
| 567 | def Repair(self, args): |
| 568 | """Repair damaged GPT headers and tables.""" |
| 569 | gpt = GPT.LoadFromFile(args.image_file) |
| 570 | gpt.Resize(os.path.getsize(args.image_file.name)) |
| 571 | |
| 572 | free_space = gpt.GetFreeSpace() |
| 573 | if args.expand: |
| 574 | if free_space: |
| 575 | gpt.ExpandPartition(0) |
| 576 | else: |
| 577 | logging.warn('- No extra space to expand.') |
| 578 | elif free_space: |
| 579 | logging.warn('Extra space found (%d, LBA=%d), ' |
Peter Shih | c7156ca | 2018-02-26 14:46:24 +0800 | [diff] [blame] | 580 | 'use --expand to expand partitions.', |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 581 | free_space, free_space / gpt.block_size) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 582 | |
| 583 | gpt.WriteToFile(args.image_file) |
| 584 | print('Disk image file %s repaired.' % args.image_file.name) |
| 585 | |
| 586 | @classmethod |
| 587 | def RegisterShow(cls, p): |
| 588 | """Registers the repair command to argparser. |
| 589 | |
| 590 | Args: |
| 591 | p: An argparse parser instance. |
| 592 | """ |
| 593 | p.add_argument('--numeric', '-n', action='store_true', |
| 594 | help='Numeric output only.') |
| 595 | p.add_argument('--quick', '-q', action='store_true', |
| 596 | help='Quick output.') |
| 597 | p.add_argument('--index', '-i', type=int, default=None, |
| 598 | help='Show specified partition only, with format args.') |
| 599 | for name, help_str in cls.FORMAT_ARGS: |
| 600 | # TODO(hungte) Alert if multiple args were specified. |
| 601 | p.add_argument('--%s' % name, '-%c' % name[0], action='store_true', |
| 602 | help='[format] %s.' % help_str) |
| 603 | p.add_argument('image_file', type=argparse.FileType('rb'), |
| 604 | help='Disk image file to show.') |
| 605 | |
| 606 | |
| 607 | def Show(self, args): |
| 608 | """Show partition table and entries.""" |
| 609 | |
| 610 | def FormatGUID(bytes_le): |
| 611 | return str(uuid.UUID(bytes_le=bytes_le)).upper() |
| 612 | |
| 613 | def FormatTypeGUID(p): |
| 614 | guid_str = FormatGUID(p.TypeGUID) |
| 615 | if not args.numeric: |
| 616 | names = gpt.TYPE_GUID_MAP.get(guid_str) |
| 617 | if names: |
| 618 | return names |
| 619 | return guid_str |
| 620 | |
| 621 | def FormatNames(p): |
| 622 | return p.Names.decode('utf-16-le').strip('\0') |
| 623 | |
| 624 | def IsBootableType(type_guid): |
| 625 | return type_guid in gpt.TYPE_GUID_LIST_BOOTABLE |
| 626 | |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 627 | def FormatAttribute(attrs): |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 628 | if args.numeric: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 629 | return '[%x]' % (attrs.raw >> 48) |
| 630 | if attrs.legacy_boot: |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 631 | return 'legacy_boot=1' |
| 632 | return 'priority=%d tries=%d successful=%d' % ( |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 633 | attrs.priority, attrs.tries, attrs.successful) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 634 | |
| 635 | def ApplyFormatArgs(p): |
| 636 | if args.begin: |
| 637 | return p.FirstLBA |
| 638 | elif args.size: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 639 | return p.blocks |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 640 | elif args.type: |
| 641 | return FormatTypeGUID(p) |
| 642 | elif args.unique: |
| 643 | return FormatGUID(p.UniqueGUID) |
| 644 | elif args.label: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 645 | return p.label |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 646 | elif args.Successful: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 647 | return p.attrs.successful |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 648 | elif args.Priority: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 649 | return p.attrs.priority |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 650 | elif args.Tries: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 651 | return p.attrs.tries |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 652 | elif args.Legacy: |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 653 | return p.attrs.legacy_boot |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 654 | elif args.Attribute: |
| 655 | return '[%x]' % (p.Attributes >> 48) |
| 656 | else: |
| 657 | return None |
| 658 | |
| 659 | def IsFormatArgsSpecified(): |
| 660 | return any(getattr(args, arg[0]) for arg in self.FORMAT_ARGS) |
| 661 | |
| 662 | gpt = GPT.LoadFromFile(args.image_file) |
| 663 | fmt = '%12s %11s %7s %s' |
| 664 | fmt2 = '%32s %s: %s' |
| 665 | header = ('start', 'size', 'part', 'contents') |
| 666 | |
| 667 | if IsFormatArgsSpecified() and args.index is None: |
| 668 | raise ValueError('Format arguments must be used with -i.') |
| 669 | |
| 670 | partitions = gpt.GetValidPartitions() |
| 671 | if not (args.index is None or 0 < args.index <= len(partitions)): |
| 672 | raise ValueError('Invalid partition index: %d' % args.index) |
| 673 | |
| 674 | do_print_gpt_blocks = False |
| 675 | if not (args.quick or IsFormatArgsSpecified()): |
| 676 | print(fmt % header) |
| 677 | if args.index is None: |
| 678 | do_print_gpt_blocks = True |
| 679 | |
| 680 | if do_print_gpt_blocks: |
| 681 | print(fmt % (gpt.header.CurrentLBA, 1, '', 'Pri GPT header')) |
| 682 | print(fmt % (gpt.header.PartitionEntriesStartingLBA, |
| 683 | gpt.GetPartitionTableBlocks(), '', 'Pri GPT table')) |
| 684 | |
| 685 | for i, p in enumerate(partitions): |
| 686 | if args.index is not None and i != args.index - 1: |
| 687 | continue |
| 688 | |
| 689 | if IsFormatArgsSpecified(): |
| 690 | print(ApplyFormatArgs(p)) |
| 691 | continue |
| 692 | |
| 693 | type_guid = FormatGUID(p.TypeGUID) |
| 694 | print(fmt % (p.FirstLBA, p.LastLBA - p.FirstLBA + 1, i + 1, |
| 695 | FormatTypeGUID(p) if args.quick else |
| 696 | 'Label: "%s"' % FormatNames(p))) |
| 697 | |
| 698 | if not args.quick: |
| 699 | print(fmt2 % ('', 'Type', FormatTypeGUID(p))) |
| 700 | print(fmt2 % ('', 'UUID', FormatGUID(p.UniqueGUID))) |
| 701 | if args.numeric or IsBootableType(type_guid): |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 702 | print(fmt2 % ('', 'Attr', FormatAttribute(p.attrs))) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 703 | |
| 704 | if do_print_gpt_blocks: |
| 705 | f = args.image_file |
Hung-Te Lin | f148d32 | 2018-04-13 10:24:42 +0800 | [diff] [blame] | 706 | f.seek(gpt.header.BackupLBA * gpt.block_size) |
Hung-Te Lin | 49ac3c2 | 2018-04-17 14:37:54 +0800 | [diff] [blame^] | 707 | backup_header = gpt.Header.ReadFrom(f) |
Hung-Te Lin | c772e1a | 2017-04-14 16:50:50 +0800 | [diff] [blame] | 708 | print(fmt % (backup_header.PartitionEntriesStartingLBA, |
| 709 | gpt.GetPartitionTableBlocks(backup_header), '', |
| 710 | 'Sec GPT table')) |
| 711 | print(fmt % (gpt.header.BackupLBA, 1, '', 'Sec GPT header')) |
| 712 | |
| 713 | def Create(self, args): |
| 714 | """Create or reset GPT headers and tables.""" |
| 715 | del args # Not used yet. |
| 716 | raise NotImplementedError |
| 717 | |
| 718 | def Add(self, args): |
| 719 | """Add, edit or remove a partition entry.""" |
| 720 | del args # Not used yet. |
| 721 | raise NotImplementedError |
| 722 | |
| 723 | def Boot(self, args): |
| 724 | """Edit the PMBR sector for legacy BIOSes.""" |
| 725 | del args # Not used yet. |
| 726 | raise NotImplementedError |
| 727 | |
| 728 | def Find(self, args): |
| 729 | """Locate a partition by its GUID.""" |
| 730 | del args # Not used yet. |
| 731 | raise NotImplementedError |
| 732 | |
| 733 | def Prioritize(self, args): |
| 734 | """Reorder the priority of all kernel partitions.""" |
| 735 | del args # Not used yet. |
| 736 | raise NotImplementedError |
| 737 | |
| 738 | def Legacy(self, args): |
| 739 | """Switch between GPT and Legacy GPT.""" |
| 740 | del args # Not used yet. |
| 741 | raise NotImplementedError |
| 742 | |
| 743 | @classmethod |
| 744 | def RegisterAllCommands(cls, subparsers): |
| 745 | """Registers all available commands to an argparser subparsers instance.""" |
| 746 | subcommands = [('show', cls.Show, cls.RegisterShow), |
| 747 | ('repair', cls.Repair, cls.RegisterRepair)] |
| 748 | for name, invocation, register_command in subcommands: |
| 749 | register_command(subparsers.add_parser(name, help=invocation.__doc__)) |
| 750 | |
| 751 | |
| 752 | def main(): |
| 753 | parser = argparse.ArgumentParser(description='GPT Utility.') |
| 754 | parser.add_argument('--verbose', '-v', action='count', default=0, |
| 755 | help='increase verbosity.') |
| 756 | parser.add_argument('--debug', '-d', action='store_true', |
| 757 | help='enable debug output.') |
| 758 | subparsers = parser.add_subparsers(help='Sub-command help.', dest='command') |
| 759 | GPTCommands.RegisterAllCommands(subparsers) |
| 760 | |
| 761 | args = parser.parse_args() |
| 762 | log_level = max(logging.WARNING - args.verbose * 10, logging.DEBUG) |
| 763 | if args.debug: |
| 764 | log_level = logging.DEBUG |
| 765 | logging.basicConfig(format='%(module)s:%(funcName)s %(message)s', |
| 766 | level=log_level) |
| 767 | commands = GPTCommands() |
| 768 | try: |
| 769 | getattr(commands, args.command.capitalize())(args) |
| 770 | except Exception as e: |
| 771 | if args.verbose or args.debug: |
| 772 | logging.exception('Failure in command [%s]', args.command) |
| 773 | exit('ERROR: %s' % e) |
| 774 | |
| 775 | |
| 776 | if __name__ == '__main__': |
| 777 | main() |