You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
364 lines
13 KiB
364 lines
13 KiB
# 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.""" |
|
|
|
from __future__ import print_function |
|
|
|
import cStringIO |
|
import hashlib |
|
import os |
|
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 = os.path.join(os.path.dirname(__file__), |
|
'payload-test-key.pem') |
|
_PUBKEY_FILE_NAME = os.path.join(os.path.dirname(__file__), |
|
'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 SetMinorVersion(self, minor_version): |
|
"""Set the payload's minor version field.""" |
|
_SetMsgField(self.manifest, 'minor_version', minor_version) |
|
|
|
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._PayloadHeader._MAGIC) |
|
_WriteInt(file_obj, payload.Payload._PayloadHeader._VERSION_SIZE, True, |
|
self.version) |
|
_WriteInt(file_obj, payload.Payload._PayloadHeader._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)
|
|
|