blob: d9cdc2a80636323af85f03472e6abf931cbf7997 [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
208class 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 Lin49ac3c22018-04-17 14:37:54 +0800215 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 Linc772e1a2017-04-14 16:50:50 +0800218 """
219
Hung-Te Linf148d322018-04-13 10:24:42 +0800220 DEFAULT_BLOCK_SIZE = 512
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800221 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 Lin49ac3c22018-04-17 14:37:54 +0800236 @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 Linc772e1a2017-04-14 16:50:50 +0800351 def __init__(self):
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800352 """GPT constructor.
353
354 See LoadFromFile for how it's usually used.
355 """
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800356 self.header = None
357 self.partitions = None
Hung-Te Linf148d322018-04-13 10:24:42 +0800358 self.block_size = self.DEFAULT_BLOCK_SIZE
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800359
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800360 @classmethod
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800361 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 Lin49ac3c22018-04-17 14:37:54 +0800371 gpt = cls()
Hung-Te Linf148d322018-04-13 10:24:42 +0800372 # Try DEFAULT_BLOCK_SIZE, then 4K.
373 for block_size in [cls.DEFAULT_BLOCK_SIZE, 4096]:
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800374 image.seek(block_size * 1)
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800375 header = gpt.Header.ReadFrom(image)
376 if header.Signature in cls.Header.SIGNATURES:
Hung-Te Linf148d322018-04-13 10:24:42 +0800377 gpt.block_size = block_size
378 break
379 else:
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800380 raise ValueError('Invalid signature in GPT header.')
Hung-Te Linf148d322018-04-13 10:24:42 +0800381
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800382 image.seek(gpt.block_size * header.PartitionEntriesStartingLBA)
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800383 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 Linc772e1a2017-04-14 16:50:50 +0800388 gpt.header = header
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800389 gpt.partitions = [
390 ReadPartition(image, i) for i in range(header.PartitionEntriesNumber)]
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800391 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 Lin49ac3c22018-04-17 14:37:54 +0800398 be zeroed. However, few implementations for example cgpt, may create
399 partitions in arbitrary order so use this carefully.
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800400 """
401 for i, p in enumerate(self.partitions):
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800402 if p.IsUnused():
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800403 return self.partitions[:i]
404 return self.partitions
405
406 def GetMaxUsedLBA(self):
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800407 """Returns the max LastLBA from all used partitions."""
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800408 parts = [p for p in self.partitions if not p.IsUnused()]
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800409 return (max(p.LastLBA for p in parts)
410 if parts else self.header.FirstUsableLBA - 1)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800411
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 Linf148d322018-04-13 10:24:42 +0800417 blocks = size / self.block_size
418 if size % self.block_size:
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800419 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 Linf148d322018-04-13 10:24:42 +0800428 old_size = self.block_size * (self.header.BackupLBA + 1)
429 if new_size % self.block_size:
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800430 raise ValueError('New file size %d is not valid for image files.' %
431 new_size)
Hung-Te Linf148d322018-04-13 10:24:42 +0800432 new_blocks = new_size / self.block_size
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800433 if old_size != new_size:
434 logging.warn('Image size (%d, LBA=%d) changed from %d (LBA=%d).',
Hung-Te Linf148d322018-04-13 10:24:42 +0800435 new_size, new_blocks, old_size, old_size / self.block_size)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800436 else:
437 logging.info('Image size (%d, LBA=%d) not changed.',
438 new_size, new_blocks)
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800439 return
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800440
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800441 # Expected location
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800442 backup_lba = new_blocks - 1
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800443 last_usable_lba = backup_lba - self.header.FirstUsableLBA
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800444
Hung-Te Lind3a2e9a2018-04-19 13:07:26 +0800445 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 Linc772e1a2017-04-14 16:50:50 +0800449
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800450 self.header = self.header.Clone(
451 BackupLBA=backup_lba, LastUsableLBA=last_usable_lba)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800452
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 Linf148d322018-04-13 10:24:42 +0800457 return self.block_size * (self.header.LastUsableLBA - max_lba)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800458
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 Lin49ac3c22018-04-17 14:37:54 +0800472 raise ValueError(
473 'Cannot expand partition %d because it is not the last allocated '
474 'partition.' % (i + 1))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800475
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800476 old_blocks = p.blocks
477 p = p.Clone(LastLBA=self.header.LastUsableLBA)
478 new_blocks = p.blocks
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800479 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 Lin49ac3c22018-04-17 14:37:54 +0800484 """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 Linc772e1a2017-04-14 16:50:50 +0800491
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 Lin49ac3c22018-04-17 14:37:54 +0800496 return self.header.Clone(
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800497 BackupLBA=self.header.CurrentLBA,
498 CurrentLBA=self.header.BackupLBA,
499 PartitionEntriesStartingLBA=partitions_starting_lba)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800500
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800501 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 Linc772e1a2017-04-14 16:50:50 +0800510
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 Linf148d322018-04-13 10:24:42 +0800514 name, lba, lba * self.block_size)
Hung-Te Lin6977ae12018-04-17 12:20:32 +0800515 image.seek(lba * self.block_size)
516 image.write(blob)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800517
518 self.UpdateChecksum()
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800519 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 Linc772e1a2017-04-14 16:50:50 +0800523 logging.info('Usable LBA: First=%d, Last=%d',
524 self.header.FirstUsableLBA, self.header.LastUsableLBA)
525 backup_header = self.GetBackupHeader()
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800526 WriteData(
527 'Backup Partitions', parts_blob,
528 backup_header.PartitionEntriesStartingLBA)
529 WriteData('Backup Header', backup_header.blob, backup_header.CurrentLBA)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800530
531
532class 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 Shihc7156ca2018-02-26 14:46:24 +0800540 ('begin', 'beginning sector'),
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800541 ('size', 'partition size (in sectors)'),
Peter Shihc7156ca2018-02-26 14:46:24 +0800542 ('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 Linc772e1a2017-04-14 16:50:50 +0800550
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 Shihc7156ca2018-02-26 14:46:24 +0800580 'use --expand to expand partitions.',
Hung-Te Linf148d322018-04-13 10:24:42 +0800581 free_space, free_space / gpt.block_size)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800582
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 Lin49ac3c22018-04-17 14:37:54 +0800627 def FormatAttribute(attrs):
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800628 if args.numeric:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800629 return '[%x]' % (attrs.raw >> 48)
630 if attrs.legacy_boot:
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800631 return 'legacy_boot=1'
632 return 'priority=%d tries=%d successful=%d' % (
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800633 attrs.priority, attrs.tries, attrs.successful)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800634
635 def ApplyFormatArgs(p):
636 if args.begin:
637 return p.FirstLBA
638 elif args.size:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800639 return p.blocks
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800640 elif args.type:
641 return FormatTypeGUID(p)
642 elif args.unique:
643 return FormatGUID(p.UniqueGUID)
644 elif args.label:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800645 return p.label
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800646 elif args.Successful:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800647 return p.attrs.successful
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800648 elif args.Priority:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800649 return p.attrs.priority
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800650 elif args.Tries:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800651 return p.attrs.tries
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800652 elif args.Legacy:
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800653 return p.attrs.legacy_boot
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800654 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 Lin49ac3c22018-04-17 14:37:54 +0800702 print(fmt2 % ('', 'Attr', FormatAttribute(p.attrs)))
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800703
704 if do_print_gpt_blocks:
705 f = args.image_file
Hung-Te Linf148d322018-04-13 10:24:42 +0800706 f.seek(gpt.header.BackupLBA * gpt.block_size)
Hung-Te Lin49ac3c22018-04-17 14:37:54 +0800707 backup_header = gpt.Header.ReadFrom(f)
Hung-Te Linc772e1a2017-04-14 16:50:50 +0800708 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
752def 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
776if __name__ == '__main__':
777 main()