| # Copyright (c) 2013 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Utilities for unit testing.""" |
| |
| import cStringIO |
| import hashlib |
| import struct |
| import subprocess |
| |
| import common |
| import payload |
| import update_metadata_pb2 |
| |
| |
| class TestError(Exception): |
| """An error during testing of update payload code.""" |
| |
| |
| # Private/public RSA keys used for testing. |
| _PRIVKEY_FILE_NAME = 'payload-test-key.pem' |
| _PUBKEY_FILE_NAME = 'payload-test-key.pub' |
| |
| |
| def KiB(count): |
| return count << 10 |
| |
| |
| def MiB(count): |
| return count << 20 |
| |
| |
| def GiB(count): |
| return count << 30 |
| |
| |
| def _WriteInt(file_obj, size, is_unsigned, val): |
| """Writes a binary-encoded integer to a file. |
| |
| It will do the correct conversion based on the reported size and whether or |
| not a signed number is expected. Assumes a network (big-endian) byte |
| ordering. |
| |
| Args: |
| file_obj: a file object |
| size: the integer size in bytes (2, 4 or 8) |
| is_unsigned: whether it is signed or not |
| val: integer value to encode |
| Raises: |
| PayloadError if a write error occurred. |
| |
| """ |
| try: |
| file_obj.write(struct.pack(common.IntPackingFmtStr(size, is_unsigned), val)) |
| except IOError, e: |
| raise payload.PayloadError('error writing to file (%s): %s' % |
| (file_obj.name, e)) |
| |
| |
| def _SetMsgField(msg, field_name, val): |
| """Sets or clears a field in a protobuf message.""" |
| if val is None: |
| msg.ClearField(field_name) |
| else: |
| setattr(msg, field_name, val) |
| |
| |
| def SignSha256(data, privkey_file_name): |
| """Signs the data's SHA256 hash with an RSA private key. |
| |
| Args: |
| data: the data whose SHA256 hash we want to sign |
| privkey_file_name: private key used for signing data |
| Returns: |
| The signature string, prepended with an ASN1 header. |
| Raises: |
| TestError if something goes wrong. |
| |
| """ |
| # pylint: disable=E1101 |
| data_sha256_hash = common.SIG_ASN1_HEADER + hashlib.sha256(data).digest() |
| sign_cmd = ['openssl', 'rsautl', '-sign', '-inkey', privkey_file_name] |
| try: |
| sign_process = subprocess.Popen(sign_cmd, stdin=subprocess.PIPE, |
| stdout=subprocess.PIPE) |
| sig, _ = sign_process.communicate(input=data_sha256_hash) |
| except Exception as e: |
| raise TestError('signing subprocess failed: %s' % e) |
| |
| return sig |
| |
| |
| class SignaturesGenerator(object): |
| """Generates a payload signatures data block.""" |
| |
| def __init__(self): |
| self.sigs = update_metadata_pb2.Signatures() |
| |
| def AddSig(self, version, data): |
| """Adds a signature to the signature sequence. |
| |
| Args: |
| version: signature version (None means do not assign) |
| data: signature binary data (None means do not assign) |
| |
| """ |
| # Pylint fails to identify a member of the Signatures message. |
| # pylint: disable=E1101 |
| sig = self.sigs.signatures.add() |
| if version is not None: |
| sig.version = version |
| if data is not None: |
| sig.data = data |
| |
| def ToBinary(self): |
| """Returns the binary representation of the signature block.""" |
| return self.sigs.SerializeToString() |
| |
| |
| class PayloadGenerator(object): |
| """Generates an update payload allowing low-level control. |
| |
| Attributes: |
| manifest: the protobuf containing the payload manifest |
| version: the payload version identifier |
| block_size: the block size pertaining to update operations |
| |
| """ |
| |
| def __init__(self, version=1): |
| self.manifest = update_metadata_pb2.DeltaArchiveManifest() |
| self.version = version |
| self.block_size = 0 |
| |
| @staticmethod |
| def _WriteExtent(ex, val): |
| """Returns an Extent message.""" |
| start_block, num_blocks = val |
| _SetMsgField(ex, 'start_block', start_block) |
| _SetMsgField(ex, 'num_blocks', num_blocks) |
| |
| @staticmethod |
| def _AddValuesToRepeatedField(repeated_field, values, write_func): |
| """Adds values to a repeated message field.""" |
| if values: |
| for val in values: |
| new_item = repeated_field.add() |
| write_func(new_item, val) |
| |
| @staticmethod |
| def _AddExtents(extents_field, values): |
| """Adds extents to an extents field.""" |
| PayloadGenerator._AddValuesToRepeatedField( |
| extents_field, values, PayloadGenerator._WriteExtent) |
| |
| def SetBlockSize(self, block_size): |
| """Sets the payload's block size.""" |
| self.block_size = block_size |
| _SetMsgField(self.manifest, 'block_size', block_size) |
| |
| def SetPartInfo(self, is_kernel, is_new, part_size, part_hash): |
| """Set the partition info entry. |
| |
| Args: |
| is_kernel: whether this is kernel partition info |
| is_new: whether to set old (False) or new (True) info |
| part_size: the partition size (in fact, filesystem size) |
| part_hash: the partition hash |
| |
| """ |
| if is_kernel: |
| # pylint: disable=E1101 |
| part_info = (self.manifest.new_kernel_info if is_new |
| else self.manifest.old_kernel_info) |
| else: |
| # pylint: disable=E1101 |
| part_info = (self.manifest.new_rootfs_info if is_new |
| else self.manifest.old_rootfs_info) |
| _SetMsgField(part_info, 'size', part_size) |
| _SetMsgField(part_info, 'hash', part_hash) |
| |
| def AddOperation(self, is_kernel, op_type, data_offset=None, |
| data_length=None, src_extents=None, src_length=None, |
| dst_extents=None, dst_length=None, data_sha256_hash=None): |
| """Adds an InstallOperation entry.""" |
| # pylint: disable=E1101 |
| operations = (self.manifest.kernel_install_operations if is_kernel |
| else self.manifest.install_operations) |
| |
| op = operations.add() |
| op.type = op_type |
| |
| _SetMsgField(op, 'data_offset', data_offset) |
| _SetMsgField(op, 'data_length', data_length) |
| |
| self._AddExtents(op.src_extents, src_extents) |
| _SetMsgField(op, 'src_length', src_length) |
| |
| self._AddExtents(op.dst_extents, dst_extents) |
| _SetMsgField(op, 'dst_length', dst_length) |
| |
| _SetMsgField(op, 'data_sha256_hash', data_sha256_hash) |
| |
| def SetSignatures(self, sigs_offset, sigs_size): |
| """Set the payload's signature block descriptors.""" |
| _SetMsgField(self.manifest, 'signatures_offset', sigs_offset) |
| _SetMsgField(self.manifest, 'signatures_size', sigs_size) |
| |
| def _WriteHeaderToFile(self, file_obj, manifest_len): |
| """Writes a payload heaer to a file.""" |
| # We need to access protected members in Payload for writing the header. |
| # pylint: disable=W0212 |
| file_obj.write(payload.Payload._MAGIC) |
| _WriteInt(file_obj, payload.Payload._VERSION_SIZE, True, self.version) |
| _WriteInt(file_obj, payload.Payload._MANIFEST_LEN_SIZE, True, manifest_len) |
| |
| def WriteToFile(self, file_obj, manifest_len=-1, data_blobs=None, |
| sigs_data=None, padding=None): |
| """Writes the payload content to a file. |
| |
| Args: |
| file_obj: a file object open for writing |
| manifest_len: manifest len to dump (otherwise computed automatically) |
| data_blobs: a list of data blobs to be concatenated to the payload |
| sigs_data: a binary Signatures message to be concatenated to the payload |
| padding: stuff to dump past the normal data blobs provided (optional) |
| |
| """ |
| manifest = self.manifest.SerializeToString() |
| if manifest_len < 0: |
| manifest_len = len(manifest) |
| self._WriteHeaderToFile(file_obj, manifest_len) |
| file_obj.write(manifest) |
| if data_blobs: |
| for data_blob in data_blobs: |
| file_obj.write(data_blob) |
| if sigs_data: |
| file_obj.write(sigs_data) |
| if padding: |
| file_obj.write(padding) |
| |
| |
| class EnhancedPayloadGenerator(PayloadGenerator): |
| """Payload generator with automatic handling of data blobs. |
| |
| Attributes: |
| data_blobs: a list of blobs, in the order they were added |
| curr_offset: the currently consumed offset of blobs added to the payload |
| |
| """ |
| |
| def __init__(self): |
| super(EnhancedPayloadGenerator, self).__init__() |
| self.data_blobs = [] |
| self.curr_offset = 0 |
| |
| def AddData(self, data_blob): |
| """Adds a (possibly orphan) data blob.""" |
| data_length = len(data_blob) |
| data_offset = self.curr_offset |
| self.curr_offset += data_length |
| self.data_blobs.append(data_blob) |
| return data_length, data_offset |
| |
| def AddOperationWithData(self, is_kernel, op_type, src_extents=None, |
| src_length=None, dst_extents=None, dst_length=None, |
| data_blob=None, do_hash_data_blob=True): |
| """Adds an install operation and associated data blob. |
| |
| This takes care of obtaining a hash of the data blob (if so instructed) |
| and appending it to the internally maintained list of blobs, including the |
| necessary offset/length accounting. |
| |
| Args: |
| is_kernel: whether this is a kernel (True) or rootfs (False) operation |
| op_type: one of REPLACE, REPLACE_BZ, MOVE or BSDIFF |
| src_extents: list of (start, length) pairs indicating src block ranges |
| src_length: size of the src data in bytes (needed for BSDIFF) |
| dst_extents: list of (start, length) pairs indicating dst block ranges |
| dst_length: size of the dst data in bytes (needed for BSDIFF) |
| data_blob: a data blob associated with this operation |
| do_hash_data_blob: whether or not to compute and add a data blob hash |
| |
| """ |
| data_offset = data_length = data_sha256_hash = None |
| if data_blob is not None: |
| if do_hash_data_blob: |
| # pylint: disable=E1101 |
| data_sha256_hash = hashlib.sha256(data_blob).digest() |
| data_length, data_offset = self.AddData(data_blob) |
| |
| self.AddOperation(is_kernel, op_type, data_offset=data_offset, |
| data_length=data_length, src_extents=src_extents, |
| src_length=src_length, dst_extents=dst_extents, |
| dst_length=dst_length, data_sha256_hash=data_sha256_hash) |
| |
| def WriteToFileWithData(self, file_obj, sigs_data=None, |
| privkey_file_name=None, |
| do_add_pseudo_operation=False, |
| is_pseudo_in_kernel=False, padding=None): |
| """Writes the payload content to a file, optionally signing the content. |
| |
| Args: |
| file_obj: a file object open for writing |
| sigs_data: signatures blob to be appended to the payload (optional; |
| payload signature fields assumed to be preset by the caller) |
| privkey_file_name: key used for signing the payload (optional; used only |
| if explicit signatures blob not provided) |
| do_add_pseudo_operation: whether a pseudo-operation should be added to |
| account for the signature blob |
| is_pseudo_in_kernel: whether the pseudo-operation should be added to |
| kernel (True) or rootfs (False) operations |
| padding: stuff to dump past the normal data blobs provided (optional) |
| Raises: |
| TestError: if arguments are inconsistent or something goes wrong. |
| |
| """ |
| sigs_len = len(sigs_data) if sigs_data else 0 |
| |
| # Do we need to generate a genuine signatures blob? |
| do_generate_sigs_data = sigs_data is None and privkey_file_name |
| |
| if do_generate_sigs_data: |
| # First, sign some arbitrary data to obtain the size of a signature blob. |
| fake_sig = SignSha256('fake-payload-data', privkey_file_name) |
| fake_sigs_gen = SignaturesGenerator() |
| fake_sigs_gen.AddSig(1, fake_sig) |
| sigs_len = len(fake_sigs_gen.ToBinary()) |
| |
| # Update the payload with proper signature attributes. |
| self.SetSignatures(self.curr_offset, sigs_len) |
| |
| # Add a pseudo-operation to account for the signature blob, if requested. |
| if do_add_pseudo_operation: |
| if not self.block_size: |
| raise TestError('cannot add pseudo-operation without knowing the ' |
| 'payload block size') |
| self.AddOperation( |
| is_pseudo_in_kernel, common.OpType.REPLACE, |
| data_offset=self.curr_offset, data_length=sigs_len, |
| dst_extents=[(common.PSEUDO_EXTENT_MARKER, |
| (sigs_len + self.block_size - 1) / self.block_size)]) |
| |
| if do_generate_sigs_data: |
| # Once all payload fields are updated, dump and sign it. |
| temp_payload_file = cStringIO.StringIO() |
| self.WriteToFile(temp_payload_file, data_blobs=self.data_blobs) |
| sig = SignSha256(temp_payload_file.getvalue(), privkey_file_name) |
| sigs_gen = SignaturesGenerator() |
| sigs_gen.AddSig(1, sig) |
| sigs_data = sigs_gen.ToBinary() |
| assert len(sigs_data) == sigs_len, 'signature blob lengths mismatch' |
| |
| # Dump the whole thing, complete with data and signature blob, to a file. |
| self.WriteToFile(file_obj, data_blobs=self.data_blobs, sigs_data=sigs_data, |
| padding=padding) |
