Add bpttool to the trout source tree am: 37dda220a3
Original change: https://android-review.googlesource.com/c/device/google/trout/+/2322591
Change-Id: I304d58d18daddbc00e202725813ba9b1ece9298d
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/tools/bpttool b/tools/bpttool
new file mode 100755
index 0000000..4b597a3
--- /dev/null
+++ b/tools/bpttool
@@ -0,0 +1,1653 @@
+#!/usr/bin/env python3
+
+# Copyright 2016, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Command-line tool for partitioning Brillo images."""
+
+
+from __future__ import absolute_import
+from __future__ import print_function
+import argparse
+import bisect
+import copy
+import functools
+import json
+import math
+import numbers
+import os
+import struct
+import sys
+import uuid
+import zlib
+from six.moves import range
+
+# Keywords used in JSON files.
+JSON_KEYWORD_SETTINGS = 'settings'
+JSON_KEYWORD_SETTINGS_AB_SUFFIXES = 'ab_suffixes'
+JSON_KEYWORD_SETTINGS_DISK_SIZE = 'disk_size'
+JSON_KEYWORD_SETTINGS_DISK_ALIGNMENT = 'disk_alignment'
+JSON_KEYWORD_SETTINGS_DISK_GUID = 'disk_guid'
+JSON_KEYWORD_SETTINGS_PARTITIONS_OFFSET_BEGIN = 'partitions_offset_begin'
+JSON_KEYWORD_PARTITIONS = 'partitions'
+JSON_KEYWORD_PARTITIONS_LABEL = 'label'
+JSON_KEYWORD_PARTITIONS_OFFSET = 'offset'
+JSON_KEYWORD_PARTITIONS_SIZE = 'size'
+JSON_KEYWORD_PARTITIONS_GROW = 'grow'
+JSON_KEYWORD_PARTITIONS_GUID = 'guid'
+JSON_KEYWORD_PARTITIONS_TYPE_GUID = 'type_guid'
+JSON_KEYWORD_PARTITIONS_FLAGS = 'flags'
+JSON_KEYWORD_PARTITIONS_PERSIST = 'persist'
+JSON_KEYWORD_PARTITIONS_IGNORE = 'ignore'
+JSON_KEYWORD_PARTITIONS_AB = 'ab'
+JSON_KEYWORD_PARTITIONS_AB_EXPANDED = 'ab_expanded'
+JSON_KEYWORD_PARTITIONS_POSITION = 'position'
+JSON_KEYWORD_AUTO = 'auto'
+
+# Possible values for the --type option of the query_partition
+# sub-command.
+QUERY_PARTITION_TYPES = ['size',
+ 'offset',
+ 'guid',
+ 'type_guid',
+ 'flags',
+ 'persist']
+
+BPT_VERSION_MAJOR = 1
+BPT_VERSION_MINOR = 0
+
+DISK_SECTOR_SIZE = 512
+
+GPT_NUM_LBAS = 33
+
+GPT_MIN_PART_NUM = 1
+GPT_MAX_PART_NUM = 128
+
+KNOWN_TYPE_GUIDS = {
+ 'brillo_boot': 'bb499290-b57e-49f6-bf41-190386693794',
+ 'brillo_bootloader': '4892aeb3-a45f-4c5f-875f-da3303c0795c',
+ 'brillo_system': '0f2778c4-5cc1-4300-8670-6c88b7e57ed6',
+ 'brillo_odm': 'e99d84d7-2c1b-44cf-8c58-effae2dc2558',
+ 'brillo_oem': 'aa3434b2-ddc3-4065-8b1a-18e99ea15cb7',
+ 'brillo_userdata': '0bb7e6ed-4424-49c0-9372-7fbab465ab4c',
+ 'brillo_misc': '6b2378b0-0fbc-4aa9-a4f6-4d6e17281c47',
+ 'brillo_vbmeta': 'b598858a-5fe3-418e-b8c4-824b41f4adfc',
+ 'brillo_vendor_specific': '314f99d5-b2bf-4883-8d03-e2f2ce507d6a',
+ 'linux_fs': '0fc63daf-8483-4772-8e79-3d69d8477de4',
+ 'ms_basic_data': 'ebd0a0a2-b9e5-4433-87c0-68b6b72699c7',
+ 'efi_system': 'c12a7328-f81f-11d2-ba4b-00a0c93ec93b'
+}
+
+
+def RoundToMultiple(number, size, round_down=False):
+ """Rounds a number up (or down) to nearest multiple of another number.
+
+ Args:
+ number: The number to round up.
+ size: The multiple to round up to.
+ round_down: If True, the number will be rounded down.
+
+ Returns:
+ If |number| is a multiple of |size|, returns |number|, otherwise
+ returns |number| + |size| - |remainder| (if |round_down| is False) or
+ |number| - |remainder| (if |round_down| is True). Always returns
+ an integer.
+ """
+ remainder = number % size
+ if remainder == 0:
+ return int(number)
+ if round_down:
+ return int(number - remainder)
+ return int(number + size - remainder)
+
+
+def ParseNumber(arg):
+ """Number parser.
+
+ If |arg| is an integer, that value is returned. Otherwise int(arg, 0)
+ is returned.
+
+ This function is suitable for use in the |type| parameter of
+ |ArgumentParser|'s add_argument() function. An improvement to just
+ using type=int is that this function supports numbers in other
+ bases, e.g. "0x1234".
+
+ Arguments:
+ arg: Argument (int or string) to parse.
+
+ Returns:
+ The parsed value, as an integer.
+
+ Raises:
+ ValueError: If the argument could not be parsed.
+ """
+ if isinstance(arg, numbers.Integral):
+ return arg
+ return int(arg, 0)
+
+
+def ParseGuid(arg):
+ """Parser for RFC 4122 GUIDs.
+
+ Arguments:
+ arg: The argument, as a string.
+
+ Returns:
+ UUID in hyphenated format.
+
+ Raises:
+ ValueError: If the given string cannot be parsed.
+ """
+ return str(uuid.UUID(arg))
+
+
+def ParseSize(arg):
+ """Parser for size strings with decimal and binary unit support.
+
+ This supports both integers and strings.
+
+ Arguments:
+ arg: The string to parse.
+
+ Returns:
+ The parsed size in bytes as an integer.
+
+ Raises:
+ ValueError: If the given string cannot be parsed.
+ """
+ if isinstance(arg, numbers.Integral):
+ return arg
+
+ ws_index = arg.find(' ')
+ if ws_index != -1:
+ num = float(arg[0:ws_index])
+ factor = 1
+ if arg.endswith('KiB'):
+ factor = 1024
+ elif arg.endswith('MiB'):
+ factor = 1024*1024
+ elif arg.endswith('GiB'):
+ factor = 1024*1024*1024
+ elif arg.endswith('TiB'):
+ factor = 1024*1024*1024*1024
+ elif arg.endswith('PiB'):
+ factor = 1024*1024*1024*1024*1024
+ elif arg.endswith('kB'):
+ factor = 1000
+ elif arg.endswith('MB'):
+ factor = 1000*1000
+ elif arg.endswith('GB'):
+ factor = 1000*1000*1000
+ elif arg.endswith('TB'):
+ factor = 1000*1000*1000*1000
+ elif arg.endswith('PB'):
+ factor = 1000*1000*1000*1000*1000
+ else:
+ raise ValueError('Cannot parse string "{}"'.format(arg))
+ value = num*factor
+ # If the resulting value isn't an integer, round up.
+ if not value.is_integer():
+ value = int(math.ceil(value))
+ else:
+ value = int(arg, 0)
+ return value
+
+
+class ImageChunk(object):
+ """Data structure used for representing chunks in Android sparse files.
+
+ Attributes:
+ chunk_type: One of TYPE_RAW, TYPE_FILL, or TYPE_DONT_CARE.
+ chunk_offset: Offset in the sparse file where this chunk begins.
+ output_offset: Offset in de-sparsified file where output begins.
+ output_size: Number of bytes in output.
+ input_offset: Offset in sparse file for data if TYPE_RAW otherwise None.
+ fill_data: Blob with data to fill if TYPE_FILL otherwise None.
+ """
+
+ FORMAT = '<2H2I'
+ TYPE_RAW = 0xcac1
+ TYPE_FILL = 0xcac2
+ TYPE_DONT_CARE = 0xcac3
+ TYPE_CRC32 = 0xcac4
+
+ def __init__(self, chunk_type, chunk_offset, output_offset, output_size,
+ input_offset, fill_data):
+ """Initializes an ImageChunk object.
+
+ Arguments:
+ chunk_type: One of TYPE_RAW, TYPE_FILL, or TYPE_DONT_CARE.
+ chunk_offset: Offset in the sparse file where this chunk begins.
+ output_offset: Offset in de-sparsified file.
+ output_size: Number of bytes in output.
+ input_offset: Offset in sparse file if TYPE_RAW otherwise None.
+ fill_data: Blob with data to fill if TYPE_FILL otherwise None.
+
+ Raises:
+ ValueError: If data is not well-formed.
+ """
+ self.chunk_type = chunk_type
+ self.chunk_offset = chunk_offset
+ self.output_offset = output_offset
+ self.output_size = output_size
+ self.input_offset = input_offset
+ self.fill_data = fill_data
+ # Check invariants.
+ if self.chunk_type == self.TYPE_RAW:
+ if self.fill_data is not None:
+ raise ValueError('RAW chunk cannot have fill_data set.')
+ if not self.input_offset:
+ raise ValueError('RAW chunk must have input_offset set.')
+ elif self.chunk_type == self.TYPE_FILL:
+ if self.fill_data is None:
+ raise ValueError('FILL chunk must have fill_data set.')
+ if self.input_offset:
+ raise ValueError('FILL chunk cannot have input_offset set.')
+ elif self.chunk_type == self.TYPE_DONT_CARE:
+ if self.fill_data is not None:
+ raise ValueError('DONT_CARE chunk cannot have fill_data set.')
+ if self.input_offset:
+ raise ValueError('DONT_CARE chunk cannot have input_offset set.')
+ else:
+ raise ValueError('Invalid chunk type')
+
+
+class ImageHandler(object):
+ """Abstraction for image I/O with support for Android sparse images.
+
+ This class provides an interface for working with image files that
+ may be using the Android Sparse Image format. When an instance is
+ constructed, we test whether it's an Android sparse file. If so,
+ operations will be on the sparse file by interpreting the sparse
+ format, otherwise they will be directly on the file. Either way the
+ operations do the same.
+
+ For reading, this interface mimics a file object - it has seek(),
+ tell(), and read() methods. For writing, only truncation
+ (truncate()) and appending is supported (append_raw(),
+ append_fill(), and append_dont_care()). Additionally, data can only
+ be written in units of the block size.
+
+ Attributes:
+ is_sparse: Whether the file being operated on is sparse.
+ block_size: The block size, typically 4096.
+ image_size: The size of the unsparsified file.
+
+ """
+ # See system/core/libsparse/sparse_format.h for details.
+ MAGIC = 0xed26ff3a
+ HEADER_FORMAT = '<I4H4I'
+
+ # These are formats and offset of just the |total_chunks| and
+ # |total_blocks| fields.
+ NUM_CHUNKS_AND_BLOCKS_FORMAT = '<II'
+ NUM_CHUNKS_AND_BLOCKS_OFFSET = 16
+
+ def __init__(self, image_filename):
+ """Initializes an image handler.
+
+ Arguments:
+ image_filename: The name of the file to operate on.
+
+ Raises:
+ ValueError: If data in the file is invalid.
+ """
+ self._image_filename = image_filename
+ self._read_header()
+
+ def _read_header(self):
+ """Initializes internal data structures used for reading file.
+
+ This may be called multiple times and is typically called after
+ modifying the file (e.g. appending, truncation).
+
+ Raises:
+ ValueError: If data in the file is invalid.
+ """
+ self.is_sparse = False
+ self.block_size = 4096
+ self._file_pos = 0
+ self._image = open(self._image_filename, 'r+b')
+ self._image.seek(0, os.SEEK_END)
+ self.image_size = self._image.tell()
+
+ self._image.seek(0, os.SEEK_SET)
+ header_bin = self._image.read(struct.calcsize(self.HEADER_FORMAT))
+ if len(header_bin) < struct.calcsize(self.HEADER_FORMAT):
+ # Not a sparse image, our job here is done.
+ return
+ (magic, major_version, minor_version, file_hdr_sz, chunk_hdr_sz,
+ block_size, self._num_total_blocks, self._num_total_chunks,
+ _) = struct.unpack(self.HEADER_FORMAT, header_bin)
+ if magic != self.MAGIC:
+ # Not a sparse image, our job here is done.
+ return
+ if not (major_version == 1 and minor_version == 0):
+ raise ValueError('Encountered sparse image format version {}.{} but '
+ 'only 1.0 is supported'.format(major_version,
+ minor_version))
+ if file_hdr_sz != struct.calcsize(self.HEADER_FORMAT):
+ raise ValueError('Unexpected file_hdr_sz value {}.'.
+ format(file_hdr_sz))
+ if chunk_hdr_sz != struct.calcsize(ImageChunk.FORMAT):
+ raise ValueError('Unexpected chunk_hdr_sz value {}.'.
+ format(chunk_hdr_sz))
+
+ self.block_size = block_size
+
+ # Build an list of chunks by parsing the file.
+ self._chunks = []
+
+ # Find the smallest offset where only "Don't care" chunks
+ # follow. This will be the size of the content in the sparse
+ # image.
+ offset = 0
+ output_offset = 0
+ for _ in range(1, self._num_total_chunks + 1):
+ chunk_offset = self._image.tell()
+
+ header_bin = self._image.read(struct.calcsize(ImageChunk.FORMAT))
+ (chunk_type, _, chunk_sz, total_sz) = struct.unpack(ImageChunk.FORMAT,
+ header_bin)
+ data_sz = total_sz - struct.calcsize(ImageChunk.FORMAT)
+
+ if chunk_type == ImageChunk.TYPE_RAW:
+ if data_sz != (chunk_sz * self.block_size):
+ raise ValueError('Raw chunk input size ({}) does not match output '
+ 'size ({})'.
+ format(data_sz, chunk_sz*self.block_size))
+ self._chunks.append(ImageChunk(ImageChunk.TYPE_RAW,
+ chunk_offset,
+ output_offset,
+ chunk_sz*self.block_size,
+ self._image.tell(),
+ None))
+ self._image.read(data_sz)
+
+ elif chunk_type == ImageChunk.TYPE_FILL:
+ if data_sz != 4:
+ raise ValueError('Fill chunk should have 4 bytes of fill, but this '
+ 'has {}'.format(data_sz))
+ fill_data = self._image.read(4)
+ self._chunks.append(ImageChunk(ImageChunk.TYPE_FILL,
+ chunk_offset,
+ output_offset,
+ chunk_sz*self.block_size,
+ None,
+ fill_data))
+ elif chunk_type == ImageChunk.TYPE_DONT_CARE:
+ if data_sz != 0:
+ raise ValueError('Don\'t care chunk input size is non-zero ({})'.
+ format(data_sz))
+ self._chunks.append(ImageChunk(ImageChunk.TYPE_DONT_CARE,
+ chunk_offset,
+ output_offset,
+ chunk_sz*self.block_size,
+ None,
+ None))
+ elif chunk_type == ImageChunk.TYPE_CRC32:
+ if data_sz != 4:
+ raise ValueError('CRC32 chunk should have 4 bytes of CRC, but '
+ 'this has {}'.format(data_sz))
+ self._image.read(4)
+ else:
+ raise ValueError('Unknown chunk type {}'.format(chunk_type))
+
+ offset += chunk_sz
+ output_offset += chunk_sz * self.block_size
+
+ # Record where sparse data end.
+ self._sparse_end = self._image.tell()
+
+ # Now that we've traversed all chunks, sanity check.
+ if self._num_total_blocks != offset:
+ raise ValueError('The header said we should have {} output blocks, '
+ 'but we saw {}'.format(self._num_total_blocks, offset))
+ junk_len = len(self._image.read())
+ if junk_len > 0:
+ raise ValueError('There were {} bytes of extra data at the end of the '
+ 'file.'.format(junk_len))
+
+ # Assign |image_size|.
+ self.image_size = output_offset
+
+ # This is used when bisecting in read() to find the initial slice.
+ self._chunk_output_offsets = [i.output_offset for i in self._chunks]
+
+ self.is_sparse = True
+
+ def _update_chunks_and_blocks(self):
+ """Helper function to update the image header.
+
+ The the |total_chunks| and |total_blocks| fields in the header
+ will be set to value of the |_num_total_blocks| and
+ |_num_total_chunks| attributes.
+
+ """
+ self._image.seek(self.NUM_CHUNKS_AND_BLOCKS_OFFSET, os.SEEK_SET)
+ self._image.write(struct.pack(self.NUM_CHUNKS_AND_BLOCKS_FORMAT,
+ self._num_total_blocks,
+ self._num_total_chunks))
+
+ def append_dont_care(self, num_bytes):
+ """Appends a DONT_CARE chunk to the sparse file.
+
+ The given number of bytes must be a multiple of the block size.
+
+ Arguments:
+ num_bytes: Size in number of bytes of the DONT_CARE chunk.
+ """
+ assert num_bytes % self.block_size == 0
+
+ if not self.is_sparse:
+ self._image.seek(0, os.SEEK_END)
+ # This is more efficient that writing NUL bytes since it'll add
+ # a hole on file systems that support sparse files (native
+ # sparse, not Android sparse).
+ self._image.truncate(self._image.tell() + num_bytes)
+ self._read_header()
+ return
+
+ self._num_total_chunks += 1
+ self._num_total_blocks += num_bytes // self.block_size
+ self._update_chunks_and_blocks()
+
+ self._image.seek(self._sparse_end, os.SEEK_SET)
+ self._image.write(struct.pack(ImageChunk.FORMAT,
+ ImageChunk.TYPE_DONT_CARE,
+ 0, # Reserved
+ num_bytes // self.block_size,
+ struct.calcsize(ImageChunk.FORMAT)))
+ self._read_header()
+
+ def append_raw(self, data):
+ """Appends a RAW chunk to the sparse file.
+
+ The length of the given data must be a multiple of the block size.
+
+ Arguments:
+ data: Data to append.
+ """
+ assert len(data) % self.block_size == 0
+
+ if not self.is_sparse:
+ self._image.seek(0, os.SEEK_END)
+ self._image.write(data)
+ self._read_header()
+ return
+
+ self._num_total_chunks += 1
+ self._num_total_blocks += len(data) // self.block_size
+ self._update_chunks_and_blocks()
+
+ self._image.seek(self._sparse_end, os.SEEK_SET)
+ self._image.write(struct.pack(ImageChunk.FORMAT,
+ ImageChunk.TYPE_RAW,
+ 0, # Reserved
+ len(data) // self.block_size,
+ len(data) +
+ struct.calcsize(ImageChunk.FORMAT)))
+ self._image.write(data)
+ self._read_header()
+
+ def append_fill(self, fill_data, size):
+ """Appends a fill chunk to the sparse file.
+
+ The total length of the fill data must be a multiple of the block size.
+
+ Arguments:
+ fill_data: Fill data to append - must be four bytes.
+ size: Number of chunk - must be a multiple of four and the block size.
+ """
+ assert len(fill_data) == 4
+ assert size % 4 == 0
+ assert size % self.block_size == 0
+
+ if not self.is_sparse:
+ self._image.seek(0, os.SEEK_END)
+ self._image.write(fill_data * (size // 4))
+ self._read_header()
+ return
+
+ self._num_total_chunks += 1
+ self._num_total_blocks += size // self.block_size
+ self._update_chunks_and_blocks()
+
+ self._image.seek(self._sparse_end, os.SEEK_SET)
+ self._image.write(struct.pack(ImageChunk.FORMAT,
+ ImageChunk.TYPE_FILL,
+ 0, # Reserved
+ size // self.block_size,
+ 4 + struct.calcsize(ImageChunk.FORMAT)))
+ self._image.write(fill_data)
+ self._read_header()
+
+ def seek(self, offset):
+ """Sets the cursor position for reading from unsparsified file.
+
+ Arguments:
+ offset: Offset to seek to from the beginning of the file.
+ """
+ self._file_pos = offset
+
+ def read(self, size):
+ """Reads data from the unsparsified file.
+
+ This method may return fewer than |size| bytes of data if the end
+ of the file was encountered.
+
+ The file cursor for reading is advanced by the number of bytes
+ read.
+
+ Arguments:
+ size: Number of bytes to read.
+
+ Returns:
+ The data.
+
+ """
+ if not self.is_sparse:
+ self._image.seek(self._file_pos)
+ data = self._image.read(size)
+ self._file_pos += len(data)
+ return data
+
+ # Iterate over all chunks.
+ chunk_idx = bisect.bisect_right(self._chunk_output_offsets,
+ self._file_pos) - 1
+ data = bytearray()
+ to_go = size
+ while to_go > 0:
+ chunk = self._chunks[chunk_idx]
+ chunk_pos_offset = self._file_pos - chunk.output_offset
+ chunk_pos_to_go = min(chunk.output_size - chunk_pos_offset, to_go)
+
+ if chunk.chunk_type == ImageChunk.TYPE_RAW:
+ self._image.seek(chunk.input_offset + chunk_pos_offset)
+ data.extend(self._image.read(chunk_pos_to_go))
+ elif chunk.chunk_type == ImageChunk.TYPE_FILL:
+ all_data = chunk.fill_data*(chunk_pos_to_go // len(chunk.fill_data) + 2)
+ offset_mod = chunk_pos_offset % len(chunk.fill_data)
+ data.extend(all_data[offset_mod:(offset_mod + chunk_pos_to_go)])
+ else:
+ assert chunk.chunk_type == ImageChunk.TYPE_DONT_CARE
+ data.extend(b'\0' * chunk_pos_to_go)
+
+ to_go -= chunk_pos_to_go
+ self._file_pos += chunk_pos_to_go
+ chunk_idx += 1
+ # Generate partial read in case of EOF.
+ if chunk_idx >= len(self._chunks):
+ break
+
+ return data
+
+ def tell(self):
+ """Returns the file cursor position for reading from unsparsified file.
+
+ Returns:
+ The file cursor position for reading.
+ """
+ return self._file_pos
+
+ def truncate(self, size):
+ """Truncates the unsparsified file.
+
+ Arguments:
+ size: Desired size of unsparsified file.
+
+ Raises:
+ ValueError: If desired size isn't a multiple of the block size.
+ """
+ if not self.is_sparse:
+ self._image.truncate(size)
+ self._read_header()
+ return
+
+ if size % self.block_size != 0:
+ raise ValueError('Cannot truncate to a size which is not a multiple '
+ 'of the block size')
+
+ if size == self.image_size:
+ # Trivial where there's nothing to do.
+ return
+ elif size < self.image_size:
+ chunk_idx = bisect.bisect_right(self._chunk_output_offsets, size) - 1
+ chunk = self._chunks[chunk_idx]
+ if chunk.output_offset != size:
+ # Truncation in the middle of a trunk - need to keep the chunk
+ # and modify it.
+ chunk_idx_for_update = chunk_idx + 1
+ num_to_keep = size - chunk.output_offset
+ assert num_to_keep % self.block_size == 0
+ if chunk.chunk_type == ImageChunk.TYPE_RAW:
+ truncate_at = (chunk.chunk_offset +
+ struct.calcsize(ImageChunk.FORMAT) + num_to_keep)
+ data_sz = num_to_keep
+ elif chunk.chunk_type == ImageChunk.TYPE_FILL:
+ truncate_at = (chunk.chunk_offset +
+ struct.calcsize(ImageChunk.FORMAT) + 4)
+ data_sz = 4
+ else:
+ assert chunk.chunk_type == ImageChunk.TYPE_DONT_CARE
+ truncate_at = chunk.chunk_offset + struct.calcsize(ImageChunk.FORMAT)
+ data_sz = 0
+ chunk_sz = num_to_keep // self.block_size
+ total_sz = data_sz + struct.calcsize(ImageChunk.FORMAT)
+ self._image.seek(chunk.chunk_offset)
+ self._image.write(struct.pack(ImageChunk.FORMAT,
+ chunk.chunk_type,
+ 0, # Reserved
+ chunk_sz,
+ total_sz))
+ chunk.output_size = num_to_keep
+ else:
+ # Truncation at trunk boundary.
+ truncate_at = chunk.chunk_offset
+ chunk_idx_for_update = chunk_idx
+
+ self._num_total_chunks = chunk_idx_for_update
+ self._num_total_blocks = 0
+ for i in range(0, chunk_idx_for_update):
+ self._num_total_blocks += self._chunks[i].output_size // self.block_size
+ self._update_chunks_and_blocks()
+ self._image.truncate(truncate_at)
+
+ # We've modified the file so re-read all data.
+ self._read_header()
+ else:
+ # Truncating to grow - just add a DONT_CARE section.
+ self.append_dont_care(size - self.image_size)
+
+
+class GuidGenerator(object):
+ """An interface for obtaining strings that are GUIDs.
+
+ To facilitate unit testing, this abstraction is used instead of the
+ directly using the uuid module.
+ """
+
+ def dispense_guid(self, partition_number):
+ """Dispenses a GUID.
+
+ Arguments:
+ partition_number: The partition number or 0 if requesting a GUID
+ for the whole disk.
+
+ Returns:
+ A RFC 4122 compliant GUID, as a string.
+ """
+ return str(uuid.uuid4())
+
+
+class Partition(object):
+ """Object representing a partition.
+
+ Attributes:
+ label: The partition label.
+ offset: Offset of the partition on the disk, or None.
+ size: Size of the partition or None if not specified.
+ grow: True if partition has been requested to use all remaining space.
+ guid: Instance GUID (RFC 4122 compliant) as a string or None or 'auto'
+ if it should be automatically generated.
+ type_guid: Type GUID (RFC 4122 compliant) as a string or a known type
+ from the |KNOWN_TYPE_GUIDS| map.
+ flags: GUID flags.
+ persist: If true, sets bit 0 of flags indicating that this partition should
+ not be deleted by the bootloader.
+ ab: If True, the partition is an A/B partition.
+ ab_expanded: If True, the A/B partitions have been generated.
+ ignore: If True, the partition should not be included in the final output.
+ position: The requested position of the partition or 0 if it doesn't matter.
+ """
+
+ def __init__(self):
+ """Initializer method."""
+ self.label = ''
+ self.offset = None
+ self.size = None
+ self.grow = False
+ self.guid = None
+ self.type_guid = None
+ self.flags = 0
+ self.persist = False
+ self.ab = False
+ self.ab_expanded = False
+ self.ignore = False
+ self.position = 0
+
+ def add_info(self, pobj):
+ """Add information to partition.
+
+ Arguments:
+ pobj: A JSON object with information about the partition.
+ """
+ self.label = pobj[JSON_KEYWORD_PARTITIONS_LABEL]
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_OFFSET)
+ if value is not None:
+ self.offset = ParseSize(value)
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_SIZE)
+ if value is not None:
+ self.size = ParseSize(value)
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_GROW)
+ if value is not None:
+ self.grow = value
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_AB)
+ if value is not None:
+ self.ab = value
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_AB_EXPANDED)
+ if value is not None:
+ self.ab_expanded = value
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_GUID)
+ if value is not None:
+ self.guid = value
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_IGNORE)
+ if value is not None:
+ self.ignore = value
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_TYPE_GUID)
+ if value is not None:
+ self.type_guid = str.lower(str(value))
+ if self.type_guid in KNOWN_TYPE_GUIDS:
+ self.type_guid = KNOWN_TYPE_GUIDS[self.type_guid]
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_FLAGS)
+ if value is not None:
+ self.flags = ParseNumber(value)
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_PERSIST)
+ if value is not None:
+ self.persist = value
+ if value:
+ self.flags = self.flags | 0x1
+ value = pobj.get(JSON_KEYWORD_PARTITIONS_POSITION)
+ if value is not None:
+ self.position = ParseNumber(value)
+
+ def expand_guid(self, guid_generator, partition_number):
+ """Assign instance GUID and type GUID if required.
+
+ Arguments:
+ guid_generator: A GuidGenerator object.
+ partition_number: The partition number, starting from 1.
+ """
+ if not self.guid or self.guid == JSON_KEYWORD_AUTO:
+ self.guid = guid_generator.dispense_guid(partition_number)
+ if not self.type_guid:
+ self.type_guid = KNOWN_TYPE_GUIDS['brillo_vendor_specific']
+
+ def validate(self):
+ """Sanity checks data in object."""
+
+ try:
+ _ = uuid.UUID(str(self.guid))
+ except ValueError:
+ raise ValueError('The string "{}" is not a valid GPT instance GUID on '
+ 'partition with label "{}".'.format(
+ str(self.guid), self.label))
+
+ try:
+ _ = uuid.UUID(str(self.type_guid))
+ except ValueError:
+ raise ValueError('The string "{}" is not a valid GPT type GUID on '
+ 'partition with label "{}".'.format(
+ str(self.type_guid), self.label))
+
+ if not self.size:
+ if not self.grow:
+ raise ValueError('Size can only be unset if "grow" is True.')
+
+ def cmp(self, other):
+ """Comparison method."""
+ self_position = self.position
+ if self_position == 0:
+ self_position = GPT_MAX_PART_NUM
+ other_position = other.position
+ if other_position == 0:
+ other_position = GPT_MAX_PART_NUM
+ if self_position < other_position:
+ return -1
+ elif self_position > other_position:
+ return 1
+ else:
+ return 0
+
+
+class Settings(object):
+ """An object for holding settings.
+
+ Attributes:
+ ab_suffixes: A list of A/B suffixes to use.
+ disk_size: An integer with the disk size in bytes.
+ partitions_offset_begin: An integer with the disk partitions
+ offset begin size in bytes.
+ disk_alignment: The alignment to use for partitions.
+ disk_guid: The GUID to use for the disk or None or 'auto' if
+ automatically generated.
+ """
+
+ def __init__(self):
+ """Initializer with defaults."""
+ self.ab_suffixes = ['_a', '_b']
+ self.disk_size = None
+ self.partitions_offset_begin = 0
+ self.disk_alignment = 4096
+ self.disk_guid = JSON_KEYWORD_AUTO
+
+
+class BptError(Exception):
+ """Application-specific errors.
+
+ These errors represent issues for which a stack-trace should not be
+ presented.
+
+ Attributes:
+ message: Error message.
+ """
+
+ def __init__(self, message):
+ Exception.__init__(self, message)
+
+
+class BptParsingError(BptError):
+ """Represents an error with an input file.
+
+ Attributes:
+ message: Error message.
+ filename: Name of the file that caused an error.
+ """
+
+ def __init__(self, filename, message):
+ self.filename = filename
+ BptError.__init__(self, message)
+
+
+class Bpt(object):
+ """Business logic for bpttool command-line tool."""
+
+ def _read_json(self, input_files, ab_collapse=True):
+ """Parses a stack of JSON files into suitable data structures.
+
+ The order of files matters as later files can modify partitions
+ declared in earlier files.
+
+ Arguments:
+ input_files: An ordered list of open files.
+ ab_collapse: If True, collapse A/B partitions.
+
+ Returns:
+ A tuple where the first element is a list of Partition objects
+ and the second element is a Settings object.
+
+ Raises:
+ BptParsingError: If an input file has an error.
+ """
+ partitions = []
+ settings = Settings()
+
+ # Read all input file and merge partitions and settings.
+ for f in input_files:
+ try:
+ obj = json.loads(f.read())
+ except ValueError as e:
+ # Unfortunately we can't easily get the line number where the
+ # error occurred.
+ raise BptParsingError(f.name, e.message)
+
+ sobj = obj.get(JSON_KEYWORD_SETTINGS)
+ if sobj:
+ ab_suffixes = sobj.get(JSON_KEYWORD_SETTINGS_AB_SUFFIXES)
+ if ab_suffixes:
+ settings.ab_suffixes = ab_suffixes
+ disk_size = sobj.get(JSON_KEYWORD_SETTINGS_DISK_SIZE)
+ if disk_size:
+ settings.disk_size = ParseSize(disk_size)
+ partitions_offset_begin = sobj.get(
+ JSON_KEYWORD_SETTINGS_PARTITIONS_OFFSET_BEGIN)
+ if partitions_offset_begin:
+ settings.partitions_offset_begin = ParseSize(partitions_offset_begin)
+ disk_alignment = sobj.get(JSON_KEYWORD_SETTINGS_DISK_ALIGNMENT)
+ if disk_alignment:
+ settings.disk_alignment = ParseSize(disk_alignment)
+ disk_guid = sobj.get(JSON_KEYWORD_SETTINGS_DISK_GUID)
+ if disk_guid:
+ settings.disk_guid = disk_guid
+
+ pobjs = obj.get(JSON_KEYWORD_PARTITIONS)
+ if pobjs:
+ for pobj in pobjs:
+ if ab_collapse and pobj.get(JSON_KEYWORD_PARTITIONS_AB_EXPANDED):
+ # If we encounter an expanded partition, unexpand it. This
+ # is to make it possible to use output-JSON (from this tool)
+ # and stack it with an input-JSON file that e.g. specifies
+ # size='256 GiB' for the 'system' partition.
+ label = pobj[JSON_KEYWORD_PARTITIONS_LABEL]
+ if label.endswith(settings.ab_suffixes[0]):
+ # Modify first A/B copy so it doesn't have the trailing suffix.
+ new_len = len(label) - len(settings.ab_suffixes[0])
+ pobj[JSON_KEYWORD_PARTITIONS_LABEL] = label[0:new_len]
+ pobj[JSON_KEYWORD_PARTITIONS_AB_EXPANDED] = False
+ pobj[JSON_KEYWORD_PARTITIONS_GUID] = JSON_KEYWORD_AUTO
+ else:
+ # Skip other A/B copies.
+ continue
+ # Find or create a partition.
+ p = None
+ for candidate in partitions:
+ if candidate.label == pobj[JSON_KEYWORD_PARTITIONS_LABEL]:
+ p = candidate
+ break
+ if not p:
+ p = Partition()
+ partitions.append(p)
+ p.add_info(pobj)
+
+ return partitions, settings
+
+ def _generate_json(self, partitions, settings):
+ """Generate a string with JSON representing partitions and settings.
+
+ Arguments:
+ partitions: A list of Partition objects.
+ settings: A Settings object.
+
+ Returns:
+ A JSON string.
+ """
+ suffixes_str = '['
+ for n in range(0, len(settings.ab_suffixes)):
+ if n != 0:
+ suffixes_str += ', '
+ suffixes_str += '"{}"'.format(settings.ab_suffixes[n])
+ suffixes_str += ']'
+
+ ret = ('{{\n'
+ ' "' + JSON_KEYWORD_SETTINGS + '": {{\n'
+ ' "' + JSON_KEYWORD_SETTINGS_AB_SUFFIXES + '": {},\n'
+ ' "' + JSON_KEYWORD_SETTINGS_PARTITIONS_OFFSET_BEGIN + '": {},\n'
+ ' "' + JSON_KEYWORD_SETTINGS_DISK_SIZE + '": {},\n'
+ ' "' + JSON_KEYWORD_SETTINGS_DISK_ALIGNMENT + '": {},\n'
+ ' "' + JSON_KEYWORD_SETTINGS_DISK_GUID + '": "{}"\n'
+ ' }},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS + '": [\n').format(
+ suffixes_str,
+ settings.partitions_offset_begin,
+ settings.disk_size,
+ settings.disk_alignment,
+ settings.disk_guid)
+
+ for n in range(0, len(partitions)):
+ p = partitions[n]
+ ret += (' {{\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_LABEL + '": "{}",\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_OFFSET + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_SIZE + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_GROW + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_GUID + '": "{}",\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_TYPE_GUID + '": "{}",\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_FLAGS + '": "{:#018x}",\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_PERSIST + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_IGNORE + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_AB + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_AB_EXPANDED + '": {},\n'
+ ' "' + JSON_KEYWORD_PARTITIONS_POSITION + '": {}\n'
+ ' }}{}\n').format(p.label,
+ p.offset,
+ p.size,
+ 'true' if p.grow else 'false',
+ p.guid,
+ p.type_guid,
+ p.flags,
+ 'true' if p.persist else 'false',
+ 'true' if p.ignore else 'false',
+ 'true' if p.ab else 'false',
+ 'true' if p.ab_expanded else 'false',
+ p.position,
+ '' if n == len(partitions) - 1 else ',')
+ ret += (' ]\n'
+ '}\n')
+ return ret
+
+ def _lba_to_chs(self, lba):
+ """Converts LBA to CHS.
+
+ Arguments:
+ lba: The sector number to convert.
+
+ Returns:
+ An array containing the CHS encoded the way it's expected in a
+ MBR partition table.
+ """
+ # See https://en.wikipedia.org/wiki/Cylinder-head-sector
+ num_heads = 255
+ num_sectors = 63
+ # If LBA isn't going to fit in CHS, return maximum CHS values.
+ max_lba = 255*num_heads*num_sectors
+ if lba > max_lba:
+ return [255, 255, 255]
+ c = lba // (num_heads*num_sectors)
+ h = (lba // num_sectors) % num_heads
+ s = lba % num_sectors
+ return [h, (((c>>8) & 0x03)<<6) | (s & 0x3f), c & 0xff]
+
+ def _generate_protective_mbr(self, settings):
+ """Generate Protective MBR.
+
+ Arguments:
+ settings: A Settings object.
+
+ Returns:
+ A string with the binary protective MBR (512 bytes).
+ """
+ # See https://en.wikipedia.org/wiki/Master_boot_record for MBR layout.
+ #
+ # The first partition starts at offset 446 (0x1be).
+ lba_start = 1
+ lba_end = settings.disk_size // DISK_SECTOR_SIZE - 1
+ start_chs = self._lba_to_chs(lba_start)
+ end_chs = self._lba_to_chs(lba_end)
+ pmbr = struct.pack('<446s' # Bootloader code
+ 'B' # Status.
+ 'BBB' # CHS start.
+ 'B' # Partition type.
+ 'BBB' # CHS end.
+ 'I' # LBA of partition start.
+ 'I' # Number of sectors in partition.
+ '48x' # Padding to get to offset 510 (0x1fe).
+ 'BB', # Boot signature.
+ b'\xfa\xeb\xfe', # cli ; jmp $ (x86)
+ 0x00,
+ start_chs[0], start_chs[1], start_chs[2],
+ 0xee, # MBR Partition Type: GPT protective MBR.
+ end_chs[0], end_chs[1], end_chs[2],
+ 1, # LBA start
+ lba_end,
+ 0x55, 0xaa)
+ return pmbr
+
+ def _generate_gpt(self, partitions, settings, primary=True):
+ """Generate GUID Partition Table.
+
+ Arguments:
+ partitions: A list of Partition objects.
+ settings: A Settings object.
+ primary: True to generate primary GPT, False to generate secondary.
+
+ Returns:
+ A string with the binary GUID Partition Table (33*512 bytes).
+ """
+ # See https://en.wikipedia.org/wiki/Master_boot_record for MBR layout.
+ #
+ # The first partition starts at offset 446 (0x1be).
+
+ disk_num_lbas = settings.disk_size // DISK_SECTOR_SIZE
+ if primary:
+ current_lba = 1
+ other_lba = disk_num_lbas - 1
+ partitions_lba = 2
+ else:
+ current_lba = disk_num_lbas - 1
+ other_lba = 1
+ partitions_lba = disk_num_lbas - GPT_NUM_LBAS
+ first_usable_lba = GPT_NUM_LBAS + 1
+ last_usable_lba = disk_num_lbas - GPT_NUM_LBAS - 1
+
+ part_array = []
+ for p in partitions:
+ part_array.append(struct.pack(
+ '<16s' # Partition type GUID.
+ '16s' # Partition instance GUID.
+ 'QQ' # First and last LBA.
+ 'Q' # Flags.
+ '72s', # Name (36 UTF-16LE code units).
+ uuid.UUID(p.type_guid).bytes_le,
+ uuid.UUID(p.guid).bytes_le,
+ p.offset // DISK_SECTOR_SIZE,
+ (p.offset + p.size) // DISK_SECTOR_SIZE - 1,
+ p.flags,
+ p.label.encode(encoding='utf-16le')))
+
+ part_array.append(((128 - len(partitions))*128) * b'\0')
+ part_array_bytes = functools.reduce(lambda x, y: x + y, part_array)
+
+ partitions_crc32 = zlib.crc32(part_array_bytes) % (1<<32)
+
+ header_crc32 = 0
+ while True:
+ header = struct.pack(
+ '<8s' # Signature.
+ '4B' # Version.
+ 'I' # Header size.
+ 'I' # CRC32 (must be zero during calculation).
+ 'I' # Reserved (must be zero).
+ 'QQ' # Current and Other LBA.
+ 'QQ' # First and last usable LBA.
+ '16s' # Disk GUID.
+ 'Q' # Starting LBA of array of partitions.
+ 'I' # Number of partitions.
+ 'I' # Partition entry size, in bytes.
+ 'I' # CRC32 of partition array
+ '420x', # Padding to get to 512 bytes.
+ b'EFI PART',
+ 0x00, 0x00, 0x01, 0x00,
+ 92,
+ header_crc32,
+ 0x00000000,
+ current_lba, other_lba,
+ first_usable_lba, last_usable_lba,
+ uuid.UUID(settings.disk_guid).bytes_le,
+ partitions_lba,
+ 128,
+ 128,
+ partitions_crc32)
+ if header_crc32 != 0:
+ break
+ header_crc32 = zlib.crc32(header[0:92]) % (1<<32)
+
+ if primary:
+ return header + part_array_bytes
+ else:
+ return part_array_bytes + header
+
+ def _generate_gpt_bin(self, partitions, settings):
+ """Generate a bytearray representing partitions and settings.
+
+ The blob will have three partition tables, laid out one after
+ another: 1) Protective MBR (512 bytes); 2) Primary GPT (33*512
+ bytes); and 3) Secondary GPT (33*512 bytes).
+
+ The total size will be 34,304 bytes.
+
+ Arguments:
+ partitions: A list of Partition objects.
+ settings: A Settings object.
+
+ Returns:
+ A bytearray() object.
+ """
+ protective_mbr = self._generate_protective_mbr(settings)
+ primary_gpt = self._generate_gpt(partitions, settings)
+ secondary_gpt = self._generate_gpt(partitions, settings, primary=False)
+ ret = protective_mbr + primary_gpt + secondary_gpt
+ return ret
+
+ def _validate_disk_partitions(self, partitions, disk_size):
+ """Check that a list of partitions have assigned offsets and fits on a
+ disk of a given size.
+
+ This function checks partition offsets and sizes to see if they may fit on
+ a disk image.
+
+ Arguments:
+ partitions: A list of Partition objects.
+ settings: Integer size of disk image.
+
+ Raises:
+ BptError: If checked condition is not satisfied.
+ """
+ for p in partitions:
+ if not p.offset or p.offset < (GPT_NUM_LBAS + 1)*DISK_SECTOR_SIZE:
+ raise BptError('Partition with label "{}" has no offset.'
+ .format(p.label))
+ if not p.size or p.size < 0:
+ raise BptError('Partition with label "{}" has no size.'
+ .format(p.label))
+ if (p.offset + p.size) > (disk_size - GPT_NUM_LBAS*DISK_SECTOR_SIZE):
+ raise BptError('Partition with label "{}" exceeds the disk '
+ 'image size.'.format(p.label))
+
+ def make_table(self,
+ inputs,
+ ab_suffixes=None,
+ partitions_offset_begin=None,
+ disk_size=None,
+ disk_alignment=None,
+ disk_guid=None,
+ guid_generator=None):
+ """Implementation of the 'make_table' command.
+
+ This function takes a list of input partition definition files,
+ flattens them, expands A/B partitions, grows partitions, and lays
+ out partitions according to alignment constraints.
+
+ Arguments:
+ inputs: List of JSON files to parse.
+ ab_suffixes: List of the A/B suffixes (as a comma-separated string)
+ to use or None to not override.
+ partitions_offset_begin: Size of disk partitions offset
+ begin or None to not override.
+ disk_size: Size of disk or None to not override.
+ disk_alignment: Disk alignment or None to not override.
+ disk_guid: Disk GUID as a string or None to not override.
+ guid_generator: A GuidGenerator or None to use the default.
+
+ Returns:
+ A tuple where the first argument is a JSON string for the resulting
+ partitions and the second argument is the binary partition tables.
+
+ Raises:
+ BptParsingError: If an input file has an error.
+ BptError: If another application-specific error occurs
+ """
+ partitions, settings = self._read_json(inputs)
+
+ # Command-line arguments override anything specified in input
+ # files.
+ if disk_size:
+ settings.disk_size = int(math.ceil(disk_size))
+ if disk_alignment:
+ settings.disk_alignment = int(disk_alignment)
+ if partitions_offset_begin:
+ settings.partitions_offset_begin = int(partitions_offset_begin)
+ if ab_suffixes:
+ settings.ab_suffixes = ab_suffixes.split(',')
+ if disk_guid:
+ settings.disk_guid = disk_guid
+
+ if not guid_generator:
+ guid_generator = GuidGenerator()
+
+ # We need to know the disk size. Also round it down to ensure it's
+ # a multiple of the sector size.
+ if not settings.disk_size:
+ raise BptError('Disk size not specified. Use --disk_size option '
+ 'or specify it in an input file.\n')
+ settings.disk_size = RoundToMultiple(settings.disk_size,
+ DISK_SECTOR_SIZE,
+ round_down=True)
+
+ # Alignment must be divisible by disk sector size.
+ if settings.disk_alignment % DISK_SECTOR_SIZE != 0:
+ raise BptError(
+ 'Disk alignment size of {} is not divisible by {}.\n'.format(
+ settings.disk_alignment, DISK_SECTOR_SIZE))
+
+ if settings.partitions_offset_begin != 0:
+ # Disk partitions offset begin size must be
+ # divisible by disk sector size.
+ if settings.partitions_offset_begin % settings.disk_alignment != 0:
+ raise BptError(
+ 'Disk Partitions offset begin size of {} '
+ 'is not divisible by {}.\n'.format(
+ settings.partitions_offset_begin, settings.disk_alignment))
+ settings.partitions_offset_begin = max(settings.partitions_offset_begin,
+ DISK_SECTOR_SIZE*(1 + GPT_NUM_LBAS))
+ settings.partitions_offset_begin = RoundToMultiple(
+ settings.partitions_offset_begin, settings.disk_alignment)
+
+ # Expand A/B partitions and skip ignored partitions.
+ expanded_partitions = []
+ for p in partitions:
+ if p.ignore:
+ continue
+ if p.ab and not p.ab_expanded:
+ p.ab_expanded = True
+ for suffix in settings.ab_suffixes:
+ new_p = copy.deepcopy(p)
+ new_p.label += suffix
+ expanded_partitions.append(new_p)
+ else:
+ expanded_partitions.append(p)
+ partitions = expanded_partitions
+
+ # Expand Disk GUID if needed.
+ if not settings.disk_guid or settings.disk_guid == JSON_KEYWORD_AUTO:
+ settings.disk_guid = guid_generator.dispense_guid(0)
+
+ # Sort according to 'position' attribute.
+ partitions = sorted(partitions, key=functools.cmp_to_key(lambda x, y: x.cmp(y)))
+
+ # Automatically generate GUIDs if the GUID is unset or set to
+ # 'auto'. Also validate the rest of the fields.
+ part_no = 1
+ for p in partitions:
+ p.expand_guid(guid_generator, part_no)
+ p.validate()
+ part_no += 1
+
+ # Idenfify partition to grow and lay out partitions, ignoring the
+ # one to grow. This way we can figure out how much space is left.
+ #
+ # Right now we only support a single 'grow' partition but we could
+ # support more in the future by splitting up the available bytes
+ # between them.
+ grow_part = None
+ # offset minimal size: DISK_SECTOR_SIZE*(1 + GPT_NUM_LBAS)
+ offset = max(settings.partitions_offset_begin,
+ DISK_SECTOR_SIZE*(1 + GPT_NUM_LBAS))
+ for p in partitions:
+ if p.grow:
+ if grow_part:
+ raise BptError('Only a single partition can be automatically '
+ 'grown.\n')
+ grow_part = p
+ else:
+ # Ensure size is a multiple of DISK_SECTOR_SIZE by rounding up
+ # (user may specify it as e.g. "1.5 GB" which is not divisible
+ # by 512).
+ p.size = RoundToMultiple(p.size, DISK_SECTOR_SIZE)
+ # Align offset to disk alignment.
+ offset = RoundToMultiple(offset, settings.disk_alignment)
+ offset += p.size
+
+ # After laying out (respecting alignment) all non-grow
+ # partitions, check that the given disk size is big enough.
+ if offset > settings.disk_size - DISK_SECTOR_SIZE*GPT_NUM_LBAS:
+ raise BptError('Disk size of {} bytes is too small for partitions '
+ 'totaling {} bytes.\n'.format(
+ settings.disk_size, offset))
+
+ # If we have a grow partition, it'll starts at the next
+ # available alignment offset and we can calculate its size as
+ # follows.
+ if grow_part:
+ offset = RoundToMultiple(offset, settings.disk_alignment)
+ grow_part.size = RoundToMultiple(
+ settings.disk_size - DISK_SECTOR_SIZE*GPT_NUM_LBAS - offset,
+ settings.disk_alignment,
+ round_down=True)
+ if grow_part.size < DISK_SECTOR_SIZE:
+ raise BptError('Not enough space for partition "{}" to be '
+ 'automatically grown.\n'.format(grow_part.label))
+
+ # Now we can assign partition start offsets for all partitions,
+ # including the grow partition.
+ # offset minimal size: DISK_SECTOR_SIZE*(1 + GPT_NUM_LBAS)
+ offset = max(settings.partitions_offset_begin,
+ DISK_SECTOR_SIZE*(1 + GPT_NUM_LBAS))
+ for p in partitions:
+ # Align offset.
+ offset = RoundToMultiple(offset, settings.disk_alignment)
+ p.offset = offset
+ offset += p.size
+ assert offset <= settings.disk_size - DISK_SECTOR_SIZE*GPT_NUM_LBAS
+
+ json_str = self._generate_json(partitions, settings)
+
+ gpt_bin = self._generate_gpt_bin(partitions, settings)
+
+ return json_str, gpt_bin
+
+ def make_disk_image(self, output, bpt, images, allow_empty_partitions=False):
+ """Implementation of the 'make_disk_image' command.
+
+ This function takes in a list of partitions images and a bpt file
+ for the purpose of creating a raw disk image with a protective MBR,
+ primary and secondary GPT, and content for each partition as specified.
+
+ Arguments:
+ output: Output file where disk image is to be written to.
+ bpt: BPT JSON file to parse.
+ images: List of partition image paths to be combined (as specified by
+ bpt). Each element is of the form.
+ 'PARTITION_NAME:/PATH/TO/PARTITION_IMAGE'
+ allow_empty_partitions: If True, partitions defined in |bpt| need not to
+ be present in |images|. Otherwise an exception is
+ thrown if a partition is referenced in |bpt| but
+ not in |images|.
+
+ Raises:
+ BptParsingError: If an image file has an error.
+ BptError: If another application-specific error occurs.
+ """
+ # Generate partition list and settings.
+ partitions, settings = self._read_json([bpt], ab_collapse=False)
+
+ # Validated partition sizes and offsets.
+ self._validate_disk_partitions(partitions, settings.disk_size)
+
+ # Sort according to 'offset' attribute.
+ partitions = sorted(partitions, key=functools.cmp_to_key(lambda x, y: x.cmp(y)))
+
+ # Create necessary tables.
+ protective_mbr = self._generate_protective_mbr(settings)
+ primary_gpt = self._generate_gpt(partitions, settings)
+ secondary_gpt = self._generate_gpt(partitions, settings, primary=False)
+
+ # Start at 0 offset for mbr and primary gpt.
+ output.seek(0)
+ output.write(protective_mbr)
+ output.write(primary_gpt)
+
+ # Create mapping of partition name to partition image file.
+ image_file_names = {}
+ try:
+ for name_path in images:
+ name, path = name_path.split(":")
+ image_file_names[name] = path
+ except ValueError as e:
+ raise BptParsingError(name_path, 'Bad image argument {}.'.format(
+ images[i]))
+
+ # Read image and insert in correct offset.
+ for p in partitions:
+ if p.label not in image_file_names:
+ if allow_empty_partitions:
+ continue
+ else:
+ raise BptParsingError(bpt.name, 'No content specified for partition'
+ ' with label {}'.format(p.label))
+
+ input_image = ImageHandler(image_file_names[p.label])
+ output.seek(p.offset)
+ partition_blob = input_image.read(p.size)
+ output.write(partition_blob)
+
+ # Put secondary GPT and end of disk.
+ output.seek(settings.disk_size - len(secondary_gpt))
+ output.write(secondary_gpt)
+
+ def query_partition(self, input_file, part_label, query_type, ab_collapse):
+ """Implementation of the 'query_partition' command.
+
+ This reads the partition definition file given by |input_file| and
+ returns information of type |query_type| for the partition with
+ label |part_label|.
+
+ Arguments:
+ input_file: A JSON file to parse.
+ part_label: Label of partition to query information about.
+ query_type: The information to query, see |QUERY_PARTITION_TYPES|.
+ ab_collapse: If True, collapse A/B partitions.
+
+ Returns:
+ The requested information as a string or None if there is no
+ partition with the given label.
+
+ Raises:
+ BptParsingError: If an input file has an error.
+ BptError: If another application-specific error occurs
+ """
+
+ partitions, _ = self._read_json([input_file], ab_collapse)
+
+ part = None
+ for p in partitions:
+ if p.label == part_label:
+ part = p
+ break
+
+ if not part:
+ return None
+
+ value = part.__dict__.get(query_type)
+ # Print out flags as a hex-value.
+ if query_type == 'flags':
+ return '{:#018x}'.format(value)
+ return str(value)
+
+
+class BptTool(object):
+ """Object for bpttool command-line tool."""
+
+ def __init__(self):
+ """Initializer method."""
+ self.bpt = Bpt()
+
+ def run(self, argv):
+ """Command-line processor.
+
+ Arguments:
+ argv: Pass sys.argv from main.
+ """
+ parser = argparse.ArgumentParser()
+ subparsers = parser.add_subparsers(title='subcommands')
+
+ sub_parser = subparsers.add_parser(
+ 'version',
+ help='Prints version of bpttool.')
+ sub_parser.set_defaults(func=self.version)
+
+ sub_parser = subparsers.add_parser(
+ 'make_table',
+ help='Lays out partitions and creates partition table.')
+ sub_parser.add_argument('--input',
+ help='Path to partition definition file.',
+ type=argparse.FileType('r'),
+ action='append')
+ sub_parser.add_argument('--ab_suffixes',
+ help='Set or override A/B suffixes.')
+ sub_parser.add_argument('--partitions_offset_begin',
+ help='Set or override disk partitions '
+ 'offset begin size.',
+ type=ParseSize)
+ sub_parser.add_argument('--disk_size',
+ help='Set or override disk size.',
+ type=ParseSize)
+ sub_parser.add_argument('--disk_alignment',
+ help='Set or override disk alignment.',
+ type=ParseSize)
+ sub_parser.add_argument('--disk_guid',
+ help='Set or override disk GUID.',
+ type=ParseGuid)
+ sub_parser.add_argument('--output_json',
+ help='JSON output file name.',
+ type=argparse.FileType('w'))
+ sub_parser.add_argument('--output_gpt',
+ help='Output file name for MBR/GPT/GPT file.',
+ type=argparse.FileType('wb'))
+ sub_parser.set_defaults(func=self.make_table)
+
+ sub_parser = subparsers.add_parser(
+ 'make_disk_image',
+ help='Creates disk image for loaded with partitions.')
+ sub_parser.add_argument('--output',
+ help='Path to image output.',
+ type=argparse.FileType('wb'),
+ required=True)
+ sub_parser.add_argument('--input',
+ help='Path to bpt file input.',
+ type=argparse.FileType('r'),
+ required=True)
+ sub_parser.add_argument('--image',
+ help='Partition name and path to image file.',
+ metavar='PARTITION_NAME:PATH',
+ action='append')
+ sub_parser.add_argument('--allow_empty_partitions',
+ help='Allow skipping partitions in bpt file.',
+ action='store_true')
+ sub_parser.set_defaults(func=self.make_disk_image)
+
+ sub_parser = subparsers.add_parser(
+ 'query_partition',
+ help='Looks up informtion about a partition.')
+ sub_parser.add_argument('--input',
+ help='Path to partition definition file.',
+ type=argparse.FileType('r'),
+ required=True)
+ sub_parser.add_argument('--label',
+ help='Label of partition to look up.',
+ required=True)
+ sub_parser.add_argument('--ab_collapse',
+ help='Collapse A/B partitions.',
+ action='store_true')
+ sub_parser.add_argument('--type',
+ help='Type of information to look up.',
+ choices=QUERY_PARTITION_TYPES,
+ required=True)
+ sub_parser.set_defaults(func=self.query_partition)
+
+ args = parser.parse_args(argv[1:])
+ args.func(args)
+
+ def version(self, _):
+ """Implements the 'version' sub-command."""
+ print('{}.{}'.format(BPT_VERSION_MAJOR, BPT_VERSION_MINOR))
+
+ def query_partition(self, args):
+ """Implements the 'query_partition' sub-command."""
+ try:
+ result = self.bpt.query_partition(args.input,
+ args.label,
+ args.type,
+ args.ab_collapse)
+ except BptParsingError as e:
+ sys.stderr.write('{}: Error parsing: {}\n'.format(e.filename, e.message))
+ sys.exit(1)
+ except BptError as e:
+ sys.stderr.write('{}\n'.format(e.message))
+ sys.exit(1)
+
+ if not result:
+ sys.stderr.write('No partition with label "{}".\n'.format(args.label))
+ sys.exit(1)
+
+ print(result)
+
+ def make_table(self, args):
+ """Implements the 'make_table' sub-command."""
+ if not args.input:
+ sys.stderr.write('Option --input is required one or more times.\n')
+ sys.exit(1)
+
+ try:
+ (json_str, gpt_bin) = self.bpt.make_table(args.input, args.ab_suffixes,
+ args.partitions_offset_begin,
+ args.disk_size,
+ args.disk_alignment,
+ args.disk_guid)
+ except BptParsingError as e:
+ sys.stderr.write('{}: Error parsing: {}\n'.format(e.filename, e.message))
+ sys.exit(1)
+ except BptError as e:
+ sys.stderr.write('{}\n'.format(e.message))
+ sys.exit(1)
+
+ if args.output_json:
+ args.output_json.write(json_str)
+ if args.output_gpt:
+ args.output_gpt.write(gpt_bin)
+
+ def make_disk_image(self, args):
+ """Implements the 'make_disk_image' sub-command."""
+ if not args.input:
+ sys.stderr.write('Option --input is required.\n')
+ sys.exit(1)
+ if not args.output:
+ sys.stderr.write('Option --ouptut is required.\n')
+ sys.exit(1)
+
+ try:
+ self.bpt.make_disk_image(args.output,
+ args.input,
+ args.image,
+ args.allow_empty_partitions)
+ except BptParsingError as e:
+ sys.stderr.write('{}: Error parsing: {}\n'.format(e.filename, e.message))
+ sys.exit(1)
+ except BptError as e:
+ sys.stderr.write('{}\n'.format(e.message))
+ sys.exit(1)
+
+if __name__ == '__main__':
+ tool = BptTool()
+ tool.run(sys.argv)
