blob: 292286dac85cb5fd3ce52c4933a7a55a414d88a7 [file] [log] [blame]
You-Cheng Syud5692942018-01-04 14:40:59 +08001#!/usr/bin/env python
Hung-Te Linc772e1a2017-04-14 16:50:50 +08002# 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
8Chromium OS factory software usually needs to access partitions from disk
9images. However, there is no good, lightweight, and portable GPT utility.
10Most Chromium OS systems use `cgpt`, but that's not by default installed on
11Ubuntu. Most systems have parted (GNU) or partx (util-linux-ng) but they have
12their own problems.
13
14For example, when a disk image is resized (usually enlarged for putting more
15resources 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
22As a result, we need a dedicated tool to help processing GPT.
23
24This pygpt.py provides a simple and customized implementation for processing
25GPT, as a replacement for `cgpt`.
26"""
27
28
29from __future__ import print_function
30
31import argparse
32import binascii
33import collections
34import logging
35import os
36import struct
37import 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 Lin49ac3c22018-04-17 14:37:54 +080043HEADER_DESCRIPTION = """
Hung-Te Linc772e1a2017-04-14 16:50:50 +080044 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 Lin49ac3c22018-04-17 14:37:54 +080061PARTITION_DESCRIPTION = """
Hung-Te Linc772e1a2017-04-14 16:50:50 +080062 16s TypeGUID
63 16s UniqueGUID
64 Q FirstLBA
65 Q LastLBA
66 Q Attributes
67 72s Names
68"""
69
70
Hung-Te Lin49ac3c22018-04-17 14:37:54 +080071def 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 Linc772e1a2017-04-14 16:50:50 +080077
78 Args:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +080079 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 Linc772e1a2017-04-14 16:50:50 +080083 """
Hung-Te Lin49ac3c22018-04-17 14:37:54 +080084 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
93class 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
116class 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 Linc772e1a2017-04-14 16:50:50 +0800206
207
Hung-Te Lin4dfd3302018-04-17 14:47:52 +0800208class GPTError(Exception):
209 """All exceptions by GPT."""
210 pass
211
212
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800213class GPT(object):
214 """A GPT helper class.
215
216 To load GPT from an existing disk image file, use `LoadFromFile`.
217 After modifications were made, use `WriteToFile` to commit changes.
218
219 Attributes:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800220 header: a namedtuple of GPT header.
221 partitions: a list of GPT partition entry nametuple.
222 block_size: integer for size of bytes in one block (sector).
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800223 """
224
Hung-Te Linf148d322018-04-13 10:24:42 +0800225 DEFAULT_BLOCK_SIZE = 512
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800226 TYPE_GUID_UNUSED = '\x00' * 16
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800227 TYPE_NAME_CHROMEOS_KERNEL = 'ChromeOS kernel'
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800228 TYPE_GUID_MAP = {
229 '00000000-0000-0000-0000-000000000000': 'Unused',
230 'EBD0A0A2-B9E5-4433-87C0-68B6B72699C7': 'Linux data',
231 'FE3A2A5D-4F32-41A7-B725-ACCC3285A309': 'ChromeOS kernel',
232 '3CB8E202-3B7E-47DD-8A3C-7FF2A13CFCEC': 'ChromeOS rootfs',
233 '2E0A753D-9E48-43B0-8337-B15192CB1B5E': 'ChromeOS reserved',
234 'CAB6E88E-ABF3-4102-A07A-D4BB9BE3C1D3': 'ChromeOS firmware',
235 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B': 'EFI System Partition',
236 }
237 TYPE_GUID_LIST_BOOTABLE = [
238 'FE3A2A5D-4F32-41A7-B725-ACCC3285A309', # ChromeOS kernel
239 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B', # EFI System Partition
240 ]
241
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800242 @GPTBlob(HEADER_DESCRIPTION)
243 class Header(GPTObject):
244 """Wrapper to Header in GPT."""
245 SIGNATURES = ['EFI PART', 'CHROMEOS']
246 SIGNATURE_IGNORE = 'IGNOREME'
247 DEFAULT_REVISION = '\x00\x00\x01\x00'
248
249 DEFAULT_PARTITION_ENTRIES = 128
250 DEFAULT_PARTITIONS_LBA = 2 # LBA 0 = MBR, LBA 1 = GPT Header.
251
252 def Clone(self, **dargs):
253 """Creates a new instance with modifications.
254
255 GPT objects are usually named tuples that are immutable, so the only way
256 to make changes is to create a new instance with modifications.
257
258 CRC32 is always updated but PartitionArrayCRC32 must be updated explicitly
259 since we can't track changes in GPT.partitions automatically.
260
261 Note since GPTHeader.Clone will always update CRC, we can only check and
262 compute CRC by super(GPT.Header, header).Clone, or header._replace.
263 """
264 dargs['CRC32'] = 0
265 header = super(GPT.Header, self).Clone(**dargs)
266 return super(GPT.Header, header).Clone(CRC32=binascii.crc32(header.blob))
267
268 class PartitionAttributes(object):
269 """Wrapper for Partition.Attributes.
270
271 This can be created using Partition.attrs, but the changed properties won't
272 apply to underlying Partition until an explicit call with
273 Partition.Clone(Attributes=new_attrs).
274 """
275
276 def __init__(self, attrs):
277 self._attrs = attrs
278
279 @property
280 def raw(self):
281 """Returns the raw integer type attributes."""
282 return self._Get()
283
284 def _Get(self):
285 return self._attrs
286
287 def _Set(self, value):
288 self._attrs = value
289
290 successful = BitProperty(_Get, _Set, 56, 1)
291 tries = BitProperty(_Get, _Set, 52, 0xf)
292 priority = BitProperty(_Get, _Set, 48, 0xf)
293 legacy_boot = BitProperty(_Get, _Set, 2, 1)
294 required = BitProperty(_Get, _Set, 0, 1)
295
296 @GPTBlob(PARTITION_DESCRIPTION)
297 class Partition(GPTObject):
298 """The partition entry in GPT.
299
300 Please include following properties when creating a Partition object:
301 - image: a string for path to the image file the partition maps to.
302 - number: the 1-based partition number.
303 - block_size: an integer for size of each block (LBA, or sector).
304 """
305 NAMES_ENCODING = 'utf-16-le'
306 NAMES_LENGTH = 72
307
308 CLONE_CONVERTERS = {
309 # TODO(hungte) check if encoded name is too long.
310 'label': lambda l: (None if l is None else
311 ('Names', l.encode(GPT.Partition.NAMES_ENCODING))),
312 'TypeGUID': lambda v: v.bytes_le if isinstance(v, uuid.UUID) else v,
313 'UniqueGUID': lambda v: v.bytes_le if isinstance(v, uuid.UUID) else v,
314 'Attributes': (
315 lambda v: v.raw if isinstance(v, GPT.PartitionAttributes) else v),
316 }
317
318 def __str__(self):
319 return '%s#%s' % (self.image, self.number)
320
321 def IsUnused(self):
322 """Returns if the partition is unused and can be allocated."""
323 return self.TypeGUID == GPT.TYPE_GUID_UNUSED
324
325 @property
326 def blocks(self):
327 """Return size of partition in blocks (see block_size)."""
328 return self.LastLBA - self.FirstLBA + 1
329
330 @property
331 def offset(self):
332 """Returns offset to partition in bytes."""
333 return self.FirstLBA * self.block_size
334
335 @property
336 def size(self):
337 """Returns size of partition in bytes."""
338 return self.blocks * self.block_size
339
340 @property
341 def type_guid(self):
342 return uuid.UUID(bytes_le=self.TypeGUID)
343
344 @property
345 def unique_guid(self):
346 return uuid.UUID(bytes_le=self.UniqueGUID)
347
348 @property
349 def label(self):
350 """Returns the Names in decoded string type."""
351 return self.Names.decode(self.NAMES_ENCODING).strip('\0')
352
353 @property
354 def attrs(self):
355 return GPT.PartitionAttributes(self.Attributes)
356
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800357 def __init__(self):
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800358 """GPT constructor.
359
360 See LoadFromFile for how it's usually used.
361 """
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800362 self.header = None
363 self.partitions = None
Hung-Te Linf148d322018-04-13 10:24:42 +0800364 self.block_size = self.DEFAULT_BLOCK_SIZE
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800365
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800366 @classmethod
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800367 def LoadFromFile(cls, image):
368 """Loads a GPT table from give disk image file object.
369
370 Args:
371 image: a string as file path or a file-like object to read from.
372 """
373 if isinstance(image, basestring):
374 with open(image, 'rb') as f:
375 return cls.LoadFromFile(f)
376
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800377 gpt = cls()
Hung-Te Linf148d322018-04-13 10:24:42 +0800378 # Try DEFAULT_BLOCK_SIZE, then 4K.
379 for block_size in [cls.DEFAULT_BLOCK_SIZE, 4096]:
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800380 image.seek(block_size * 1)
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800381 header = gpt.Header.ReadFrom(image)
382 if header.Signature in cls.Header.SIGNATURES:
Hung-Te Linf148d322018-04-13 10:24:42 +0800383 gpt.block_size = block_size
384 break
385 else:
Hung-Te Lin4dfd3302018-04-17 14:47:52 +0800386 raise GPTError('Invalid signature in GPT header.')
Hung-Te Linf148d322018-04-13 10:24:42 +0800387
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800388 image.seek(gpt.block_size * header.PartitionEntriesStartingLBA)
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800389 def ReadPartition(image, i):
390 p = gpt.Partition.ReadFrom(
391 image, image=image.name, number=i + 1, block_size=gpt.block_size)
392 return p
393
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800394 gpt.header = header
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800395 gpt.partitions = [
396 ReadPartition(image, i) for i in range(header.PartitionEntriesNumber)]
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800397 return gpt
398
399 def GetValidPartitions(self):
400 """Returns the list of partitions before entry with empty type GUID.
401
402 In partition table, the first entry with empty type GUID indicates end of
403 valid partitions. In most implementations all partitions after that should
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800404 be zeroed. However, few implementations for example cgpt, may create
405 partitions in arbitrary order so use this carefully.
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800406 """
407 for i, p in enumerate(self.partitions):
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800408 if p.IsUnused():
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800409 return self.partitions[:i]
410 return self.partitions
411
412 def GetMaxUsedLBA(self):
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800413 """Returns the max LastLBA from all used partitions."""
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800414 parts = [p for p in self.partitions if not p.IsUnused()]
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800415 return (max(p.LastLBA for p in parts)
416 if parts else self.header.FirstUsableLBA - 1)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800417
418 def GetPartitionTableBlocks(self, header=None):
419 """Returns the blocks (or LBA) of partition table from given header."""
420 if header is None:
421 header = self.header
422 size = header.PartitionEntrySize * header.PartitionEntriesNumber
Hung-Te Linf148d322018-04-13 10:24:42 +0800423 blocks = size / self.block_size
424 if size % self.block_size:
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800425 blocks += 1
426 return blocks
427
428 def Resize(self, new_size):
429 """Adjust GPT for a disk image in given size.
430
431 Args:
432 new_size: Integer for new size of disk image file.
433 """
Hung-Te Linf148d322018-04-13 10:24:42 +0800434 old_size = self.block_size * (self.header.BackupLBA + 1)
435 if new_size % self.block_size:
Hung-Te Lin4dfd3302018-04-17 14:47:52 +0800436 raise GPTError(
437 'New file size %d is not valid for image files.' % new_size)
Hung-Te Linf148d322018-04-13 10:24:42 +0800438 new_blocks = new_size / self.block_size
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800439 if old_size != new_size:
440 logging.warn('Image size (%d, LBA=%d) changed from %d (LBA=%d).',
Hung-Te Linf148d322018-04-13 10:24:42 +0800441 new_size, new_blocks, old_size, old_size / self.block_size)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800442 else:
443 logging.info('Image size (%d, LBA=%d) not changed.',
444 new_size, new_blocks)
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800445 return
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800446
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800447 # Expected location
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800448 backup_lba = new_blocks - 1
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800449 last_usable_lba = backup_lba - self.header.FirstUsableLBA
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800450
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800451 if last_usable_lba < self.header.LastUsableLBA:
452 max_used_lba = self.GetMaxUsedLBA()
453 if last_usable_lba < max_used_lba:
Hung-Te Lin4dfd3302018-04-17 14:47:52 +0800454 raise GPTError('Backup partition tables will overlap used partitions')
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800455
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800456 self.header = self.header.Clone(
457 BackupLBA=backup_lba, LastUsableLBA=last_usable_lba)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800458
459 def GetFreeSpace(self):
460 """Returns the free (available) space left according to LastUsableLBA."""
461 max_lba = self.GetMaxUsedLBA()
462 assert max_lba <= self.header.LastUsableLBA, "Partitions too large."
Hung-Te Linf148d322018-04-13 10:24:42 +0800463 return self.block_size * (self.header.LastUsableLBA - max_lba)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800464
465 def ExpandPartition(self, i):
466 """Expands a given partition to last usable LBA.
467
468 Args:
469 i: Index (0-based) of target partition.
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800470
471 Returns:
472 (old_blocks, new_blocks) for size in blocks.
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800473 """
474 # Assume no partitions overlap, we need to make sure partition[i] has
475 # largest LBA.
476 if i < 0 or i >= len(self.GetValidPartitions()):
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800477 raise GPTError('Partition number %d is invalid.' % (i + 1))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800478 p = self.partitions[i]
479 max_used_lba = self.GetMaxUsedLBA()
480 if max_used_lba > p.LastLBA:
Hung-Te Lin4dfd3302018-04-17 14:47:52 +0800481 raise GPTError(
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800482 'Cannot expand partition %d because it is not the last allocated '
483 'partition.' % (i + 1))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800484
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800485 old_blocks = p.blocks
486 p = p.Clone(LastLBA=self.header.LastUsableLBA)
487 new_blocks = p.blocks
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800488 self.partitions[i] = p
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800489 logging.warn(
490 '%s expanded, size in LBA: %d -> %d.', p, old_blocks, new_blocks)
491 return (old_blocks, new_blocks)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800492
493 def UpdateChecksum(self):
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800494 """Updates all checksum fields in GPT objects.
495
496 The Header.CRC32 is automatically updated in Header.Clone().
497 """
498 parts = ''.join(p.blob for p in self.partitions)
499 self.header = self.header.Clone(
500 PartitionArrayCRC32=binascii.crc32(parts))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800501
502 def GetBackupHeader(self):
503 """Returns the backup header according to current header."""
504 partitions_starting_lba = (
505 self.header.BackupLBA - self.GetPartitionTableBlocks())
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800506 return self.header.Clone(
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800507 BackupLBA=self.header.CurrentLBA,
508 CurrentLBA=self.header.BackupLBA,
509 PartitionEntriesStartingLBA=partitions_starting_lba)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800510
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800511 def WriteToFile(self, image):
512 """Updates partition table in a disk image file.
513
514 Args:
515 image: a string as file path or a file-like object to write into.
516 """
517 if isinstance(image, basestring):
518 with open(image, 'rb+') as f:
519 return self.WriteToFile(f)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800520
521 def WriteData(name, blob, lba):
522 """Writes a blob into given location."""
523 logging.info('Writing %s in LBA %d (offset %d)',
Hung-Te Linf148d322018-04-13 10:24:42 +0800524 name, lba, lba * self.block_size)
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800525 image.seek(lba * self.block_size)
526 image.write(blob)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800527
528 self.UpdateChecksum()
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800529 parts_blob = ''.join(p.blob for p in self.partitions)
530 WriteData('GPT Header', self.header.blob, self.header.CurrentLBA)
531 WriteData(
532 'GPT Partitions', parts_blob, self.header.PartitionEntriesStartingLBA)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800533 logging.info('Usable LBA: First=%d, Last=%d',
534 self.header.FirstUsableLBA, self.header.LastUsableLBA)
535 backup_header = self.GetBackupHeader()
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800536 WriteData(
537 'Backup Partitions', parts_blob,
538 backup_header.PartitionEntriesStartingLBA)
539 WriteData('Backup Header', backup_header.blob, backup_header.CurrentLBA)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800540
541
542class GPTCommands(object):
543 """Collection of GPT sub commands for command line to use.
544
545 The commands are derived from `cgpt`, but not necessary to be 100% compatible
546 with cgpt.
547 """
548
549 FORMAT_ARGS = [
Peter Shihc7156ca2018-02-26 14:46:24 +0800550 ('begin', 'beginning sector'),
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800551 ('size', 'partition size (in sectors)'),
Peter Shihc7156ca2018-02-26 14:46:24 +0800552 ('type', 'type guid'),
553 ('unique', 'unique guid'),
554 ('label', 'label'),
555 ('Successful', 'Successful flag'),
556 ('Tries', 'Tries flag'),
557 ('Priority', 'Priority flag'),
558 ('Legacy', 'Legacy Boot flag'),
559 ('Attribute', 'raw 16-bit attribute value (bits 48-63)')]
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800560
561 def __init__(self):
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800562 commands = dict(
563 (command.lower(), getattr(self, command)())
564 for command in dir(self)
565 if (isinstance(getattr(self, command), type) and
566 issubclass(getattr(self, command), self.SubCommand) and
567 getattr(self, command) is not self.SubCommand)
568 )
569 self.commands = commands
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800570
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800571 def DefineArgs(self, parser):
572 """Defines all available commands to an argparser subparsers instance."""
573 subparsers = parser.add_subparsers(help='Sub-command help.', dest='command')
574 for name, instance in sorted(self.commands.iteritems()):
575 parser = subparsers.add_parser(
576 name, description=instance.__doc__,
577 formatter_class=argparse.RawDescriptionHelpFormatter,
578 help=instance.__doc__.splitlines()[0])
579 instance.DefineArgs(parser)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800580
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800581 def Execute(self, args):
582 """Execute the sub commands by given parsed arguments."""
583 self.commands[args.command].Execute(args)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800584
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800585 class SubCommand(object):
586 """A base class for sub commands to derive from."""
587
588 def DefineArgs(self, parser):
589 """Defines command line arguments to argparse parser.
590
591 Args:
592 parser: An argparse parser instance.
593 """
594 del parser # Unused.
595 raise NotImplementedError
596
597 def Execute(self, args):
598 """Execute the command.
599
600 Args:
601 args: An argparse parsed namespace.
602 """
603 del args # Unused.
604 raise NotImplementedError
605
606 class Repair(SubCommand):
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800607 """Repair damaged GPT headers and tables."""
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800608
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800609 def DefineArgs(self, parser):
610 parser.add_argument(
611 'image_file', type=argparse.FileType('rb+'),
612 help='Disk image file to repair.')
613
614 def Execute(self, args):
615 gpt = GPT.LoadFromFile(args.image_file)
616 gpt.Resize(os.path.getsize(args.image_file.name))
617 gpt.WriteToFile(args.image_file)
618 print('Disk image file %s repaired.' % args.image_file.name)
619
620 class Expand(SubCommand):
621 """Expands a GPT partition to all available free space."""
622
623 def DefineArgs(self, parser):
624 parser.add_argument(
625 '-i', '--number', type=int, required=True,
626 help='The partition to expand.')
627 parser.add_argument(
628 'image_file', type=argparse.FileType('rb+'),
629 help='Disk image file to modify.')
630
631 def Execute(self, args):
632 gpt = GPT.LoadFromFile(args.image_file)
633 old_blocks, new_blocks = gpt.ExpandPartition(args.number - 1)
634 gpt.WriteToFile(args.image_file)
635 if old_blocks < new_blocks:
636 print(
637 'Partition %s on disk image file %s has been extended '
638 'from %s to %s .' %
639 (args.number, args.image_file.name, old_blocks * gpt.block_size,
640 new_blocks * gpt.block_size))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800641 else:
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800642 print('Nothing to expand for disk image %s partition %s.' %
643 (args.image_file.name, args.number))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800644
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800645 class Show(SubCommand):
646 """Show partition table and entries.
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800647
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800648 Display the GPT table.
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800649 """
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800650
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800651 def DefineArgs(self, parser):
652 parser.add_argument(
653 '--numeric', '-n', action='store_true',
654 help='Numeric output only.')
655 parser.add_argument(
656 '--quick', '-q', action='store_true',
657 help='Quick output.')
658 parser.add_argument(
659 '-i', '--number', type=int,
660 help='Show specified partition only, with format args.')
661 for name, help_str in GPTCommands.FORMAT_ARGS:
662 # TODO(hungte) Alert if multiple args were specified.
663 parser.add_argument(
664 '--%s' % name, '-%c' % name[0], action='store_true',
665 help='[format] %s.' % help_str)
666 parser.add_argument(
667 'image_file', type=argparse.FileType('rb'),
668 help='Disk image file to show.')
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800669
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800670 def Execute(self, args):
671 """Show partition table and entries."""
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800672
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800673 def FormatGUID(bytes_le):
674 return str(uuid.UUID(bytes_le=bytes_le)).upper()
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800675
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800676 def FormatTypeGUID(p):
677 guid_str = FormatGUID(p.TypeGUID)
678 if not args.numeric:
679 names = gpt.TYPE_GUID_MAP.get(guid_str)
680 if names:
681 return names
682 return guid_str
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800683
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800684 def FormatNames(p):
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800685 return p.label
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800686
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800687 def IsBootableType(type_guid):
688 return type_guid in gpt.TYPE_GUID_LIST_BOOTABLE
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800689
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800690 def FormatAttribute(attrs, chromeos_kernel=False):
691 if args.numeric:
692 return '[%x]' % (attrs.raw >> 48)
693 results = []
694 if chromeos_kernel:
695 results += [
696 'priority=%d' % attrs.priority,
697 'tries=%d' % attrs.tries,
698 'successful=%d' % attrs.successful]
699 if attrs.required:
700 results += ['required=1']
701 if attrs.legacy_boot:
702 results += ['legacy_boot=1']
703 return ' '.join(results)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800704
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800705 def ApplyFormatArgs(p):
706 if args.begin:
707 return p.FirstLBA
708 elif args.size:
709 return p.blocks
710 elif args.type:
711 return FormatTypeGUID(p)
712 elif args.unique:
713 return FormatGUID(p.UniqueGUID)
714 elif args.label:
715 return FormatNames(p)
716 elif args.Successful:
717 return p.attrs.successful
718 elif args.Priority:
719 return p.attrs.priority
720 elif args.Tries:
721 return p.attrs.tries
722 elif args.Legacy:
723 return p.attrs.legacy_boot
724 elif args.Attribute:
725 return '[%x]' % (p.Attributes >> 48)
726 else:
727 return None
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800728
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800729 def IsFormatArgsSpecified():
730 return any(getattr(args, arg[0]) for arg in GPTCommands.FORMAT_ARGS)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800731
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800732 gpt = GPT.LoadFromFile(args.image_file)
733 logging.debug('%r', gpt.header)
734 fmt = '%12s %11s %7s %s'
735 fmt2 = '%32s %s: %s'
736 header = ('start', 'size', 'part', 'contents')
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800737
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800738 if IsFormatArgsSpecified() and args.number is None:
739 raise GPTError('Format arguments must be used with -i.')
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800740
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800741 if not (args.number is None or
742 0 < args.number <= gpt.header.PartitionEntriesNumber):
743 raise GPTError('Invalid partition number: %d' % args.number)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800744
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800745 partitions = gpt.partitions
746 do_print_gpt_blocks = False
747 if not (args.quick or IsFormatArgsSpecified()):
748 print(fmt % header)
749 if args.number is None:
750 do_print_gpt_blocks = True
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800751
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800752 if do_print_gpt_blocks:
753 print(fmt % (gpt.header.CurrentLBA, 1, '', 'Pri GPT header'))
754 print(fmt % (gpt.header.PartitionEntriesStartingLBA,
755 gpt.GetPartitionTableBlocks(), '', 'Pri GPT table'))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800756
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800757 for p in partitions:
758 if args.number is None:
759 # Skip unused partitions.
760 if p.IsUnused():
761 continue
762 elif p.number != args.number:
763 continue
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800764
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800765 if IsFormatArgsSpecified():
766 print(ApplyFormatArgs(p))
767 continue
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800768
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800769 type_guid = FormatGUID(p.TypeGUID)
770 print(fmt % (p.FirstLBA, p.blocks, p.number,
771 FormatTypeGUID(p) if args.quick else
772 'Label: "%s"' % FormatNames(p)))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800773
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800774 if not args.quick:
775 print(fmt2 % ('', 'Type', FormatTypeGUID(p)))
776 print(fmt2 % ('', 'UUID', FormatGUID(p.UniqueGUID)))
777 if args.numeric or IsBootableType(type_guid):
778 name = GPT.TYPE_GUID_MAP[type_guid]
779 print(fmt2 % ('', 'Attr', FormatAttribute(
780 p.attrs, name == GPT.TYPE_NAME_CHROMEOS_KERNEL)))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800781
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800782 if do_print_gpt_blocks:
783 f = args.image_file
784 f.seek(gpt.header.BackupLBA * gpt.block_size)
785 header = gpt.Header.ReadFrom(f)
786 print(fmt % (header.PartitionEntriesStartingLBA,
787 gpt.GetPartitionTableBlocks(header), '',
788 'Sec GPT table'))
789 print(fmt % (header.CurrentLBA, 1, '', 'Sec GPT header'))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800790
791
792def main():
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800793 commands = GPTCommands()
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800794 parser = argparse.ArgumentParser(description='GPT Utility.')
795 parser.add_argument('--verbose', '-v', action='count', default=0,
796 help='increase verbosity.')
797 parser.add_argument('--debug', '-d', action='store_true',
798 help='enable debug output.')
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800799 commands.DefineArgs(parser)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800800
801 args = parser.parse_args()
802 log_level = max(logging.WARNING - args.verbose * 10, logging.DEBUG)
803 if args.debug:
804 log_level = logging.DEBUG
805 logging.basicConfig(format='%(module)s:%(funcName)s %(message)s',
806 level=log_level)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800807 try:
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800808 commands.Execute(args)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800809 except Exception as e:
810 if args.verbose or args.debug:
811 logging.exception('Failure in command [%s]', args.command)
Hung-Te Lin5cb0c312018-04-17 14:56:43 +0800812 exit('ERROR: %s: %s' % (args.command, str(e) or 'Unknown error.'))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800813
814
815if __name__ == '__main__':
816 main()