1# 2# Copyright (C) 2013 The Android Open Source Project 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15# 16 17"""Utilities for unit testing.""" 18 19from __future__ import print_function 20 21import cStringIO 22import hashlib 23import os 24import struct 25import subprocess 26 27from update_payload import common 28from update_payload import payload 29from update_payload import update_metadata_pb2 30 31 32class TestError(Exception): 33 """An error during testing of update payload code.""" 34 35 36# Private/public RSA keys used for testing. 37_PRIVKEY_FILE_NAME = os.path.join(os.path.dirname(__file__), 38 'payload-test-key.pem') 39_PUBKEY_FILE_NAME = os.path.join(os.path.dirname(__file__), 40 'payload-test-key.pub') 41 42 43def KiB(count): 44 return count << 10 45 46 47def MiB(count): 48 return count << 20 49 50 51def GiB(count): 52 return count << 30 53 54 55def _WriteInt(file_obj, size, is_unsigned, val): 56 """Writes a binary-encoded integer to a file. 57 58 It will do the correct conversion based on the reported size and whether or 59 not a signed number is expected. Assumes a network (big-endian) byte 60 ordering. 61 62 Args: 63 file_obj: a file object 64 size: the integer size in bytes (2, 4 or 8) 65 is_unsigned: whether it is signed or not 66 val: integer value to encode 67 68 Raises: 69 PayloadError if a write error occurred. 70 """ 71 try: 72 file_obj.write(struct.pack(common.IntPackingFmtStr(size, is_unsigned), val)) 73 except IOError, e: 74 raise payload.PayloadError('error writing to file (%s): %s' % 75 (file_obj.name, e)) 76 77 78def _SetMsgField(msg, field_name, val): 79 """Sets or clears a field in a protobuf message.""" 80 if val is None: 81 msg.ClearField(field_name) 82 else: 83 setattr(msg, field_name, val) 84 85 86def SignSha256(data, privkey_file_name): 87 """Signs the data's SHA256 hash with an RSA private key. 88 89 Args: 90 data: the data whose SHA256 hash we want to sign 91 privkey_file_name: private key used for signing data 92 93 Returns: 94 The signature string, prepended with an ASN1 header. 95 96 Raises: 97 TestError if something goes wrong. 98 """ 99 data_sha256_hash = common.SIG_ASN1_HEADER + hashlib.sha256(data).digest() 100 sign_cmd = ['openssl', 'rsautl', '-sign', '-inkey', privkey_file_name] 101 try: 102 sign_process = subprocess.Popen(sign_cmd, stdin=subprocess.PIPE, 103 stdout=subprocess.PIPE) 104 sig, _ = sign_process.communicate(input=data_sha256_hash) 105 except Exception as e: 106 raise TestError('signing subprocess failed: %s' % e) 107 108 return sig 109 110 111class SignaturesGenerator(object): 112 """Generates a payload signatures data block.""" 113 114 def __init__(self): 115 self.sigs = update_metadata_pb2.Signatures() 116 117 def AddSig(self, version, data): 118 """Adds a signature to the signature sequence. 119 120 Args: 121 version: signature version (None means do not assign) 122 data: signature binary data (None means do not assign) 123 """ 124 sig = self.sigs.signatures.add() 125 if version is not None: 126 sig.version = version 127 if data is not None: 128 sig.data = data 129 130 def ToBinary(self): 131 """Returns the binary representation of the signature block.""" 132 return self.sigs.SerializeToString() 133 134 135class PayloadGenerator(object): 136 """Generates an update payload allowing low-level control. 137 138 Attributes: 139 manifest: the protobuf containing the payload manifest 140 version: the payload version identifier 141 block_size: the block size pertaining to update operations 142 143 """ 144 145 def __init__(self, version=1): 146 self.manifest = update_metadata_pb2.DeltaArchiveManifest() 147 self.version = version 148 self.block_size = 0 149 150 @staticmethod 151 def _WriteExtent(ex, val): 152 """Returns an Extent message.""" 153 start_block, num_blocks = val 154 _SetMsgField(ex, 'start_block', start_block) 155 _SetMsgField(ex, 'num_blocks', num_blocks) 156 157 @staticmethod 158 def _AddValuesToRepeatedField(repeated_field, values, write_func): 159 """Adds values to a repeated message field.""" 160 if values: 161 for val in values: 162 new_item = repeated_field.add() 163 write_func(new_item, val) 164 165 @staticmethod 166 def _AddExtents(extents_field, values): 167 """Adds extents to an extents field.""" 168 PayloadGenerator._AddValuesToRepeatedField( 169 extents_field, values, PayloadGenerator._WriteExtent) 170 171 def SetBlockSize(self, block_size): 172 """Sets the payload's block size.""" 173 self.block_size = block_size 174 _SetMsgField(self.manifest, 'block_size', block_size) 175 176 def SetPartInfo(self, is_kernel, is_new, part_size, part_hash): 177 """Set the partition info entry. 178 179 Args: 180 is_kernel: whether this is kernel partition info 181 is_new: whether to set old (False) or new (True) info 182 part_size: the partition size (in fact, filesystem size) 183 part_hash: the partition hash 184 """ 185 if is_kernel: 186 part_info = (self.manifest.new_kernel_info if is_new 187 else self.manifest.old_kernel_info) 188 else: 189 part_info = (self.manifest.new_rootfs_info if is_new 190 else self.manifest.old_rootfs_info) 191 _SetMsgField(part_info, 'size', part_size) 192 _SetMsgField(part_info, 'hash', part_hash) 193 194 def AddOperation(self, is_kernel, op_type, data_offset=None, 195 data_length=None, src_extents=None, src_length=None, 196 dst_extents=None, dst_length=None, data_sha256_hash=None): 197 """Adds an InstallOperation entry.""" 198 operations = (self.manifest.kernel_install_operations if is_kernel 199 else self.manifest.install_operations) 200 201 op = operations.add() 202 op.type = op_type 203 204 _SetMsgField(op, 'data_offset', data_offset) 205 _SetMsgField(op, 'data_length', data_length) 206 207 self._AddExtents(op.src_extents, src_extents) 208 _SetMsgField(op, 'src_length', src_length) 209 210 self._AddExtents(op.dst_extents, dst_extents) 211 _SetMsgField(op, 'dst_length', dst_length) 212 213 _SetMsgField(op, 'data_sha256_hash', data_sha256_hash) 214 215 def SetSignatures(self, sigs_offset, sigs_size): 216 """Set the payload's signature block descriptors.""" 217 _SetMsgField(self.manifest, 'signatures_offset', sigs_offset) 218 _SetMsgField(self.manifest, 'signatures_size', sigs_size) 219 220 def SetMinorVersion(self, minor_version): 221 """Set the payload's minor version field.""" 222 _SetMsgField(self.manifest, 'minor_version', minor_version) 223 224 def _WriteHeaderToFile(self, file_obj, manifest_len): 225 """Writes a payload heaer to a file.""" 226 # We need to access protected members in Payload for writing the header. 227 # pylint: disable=W0212 228 file_obj.write(payload.Payload._PayloadHeader._MAGIC) 229 _WriteInt(file_obj, payload.Payload._PayloadHeader._VERSION_SIZE, True, 230 self.version) 231 _WriteInt(file_obj, payload.Payload._PayloadHeader._MANIFEST_LEN_SIZE, True, 232 manifest_len) 233 234 def WriteToFile(self, file_obj, manifest_len=-1, data_blobs=None, 235 sigs_data=None, padding=None): 236 """Writes the payload content to a file. 237 238 Args: 239 file_obj: a file object open for writing 240 manifest_len: manifest len to dump (otherwise computed automatically) 241 data_blobs: a list of data blobs to be concatenated to the payload 242 sigs_data: a binary Signatures message to be concatenated to the payload 243 padding: stuff to dump past the normal data blobs provided (optional) 244 """ 245 manifest = self.manifest.SerializeToString() 246 if manifest_len < 0: 247 manifest_len = len(manifest) 248 self._WriteHeaderToFile(file_obj, manifest_len) 249 file_obj.write(manifest) 250 if data_blobs: 251 for data_blob in data_blobs: 252 file_obj.write(data_blob) 253 if sigs_data: 254 file_obj.write(sigs_data) 255 if padding: 256 file_obj.write(padding) 257 258 259class EnhancedPayloadGenerator(PayloadGenerator): 260 """Payload generator with automatic handling of data blobs. 261 262 Attributes: 263 data_blobs: a list of blobs, in the order they were added 264 curr_offset: the currently consumed offset of blobs added to the payload 265 """ 266 267 def __init__(self): 268 super(EnhancedPayloadGenerator, self).__init__() 269 self.data_blobs = [] 270 self.curr_offset = 0 271 272 def AddData(self, data_blob): 273 """Adds a (possibly orphan) data blob.""" 274 data_length = len(data_blob) 275 data_offset = self.curr_offset 276 self.curr_offset += data_length 277 self.data_blobs.append(data_blob) 278 return data_length, data_offset 279 280 def AddOperationWithData(self, is_kernel, op_type, src_extents=None, 281 src_length=None, dst_extents=None, dst_length=None, 282 data_blob=None, do_hash_data_blob=True): 283 """Adds an install operation and associated data blob. 284 285 This takes care of obtaining a hash of the data blob (if so instructed) 286 and appending it to the internally maintained list of blobs, including the 287 necessary offset/length accounting. 288 289 Args: 290 is_kernel: whether this is a kernel (True) or rootfs (False) operation 291 op_type: one of REPLACE, REPLACE_BZ, REPLACE_XZ, MOVE or BSDIFF 292 src_extents: list of (start, length) pairs indicating src block ranges 293 src_length: size of the src data in bytes (needed for BSDIFF) 294 dst_extents: list of (start, length) pairs indicating dst block ranges 295 dst_length: size of the dst data in bytes (needed for BSDIFF) 296 data_blob: a data blob associated with this operation 297 do_hash_data_blob: whether or not to compute and add a data blob hash 298 """ 299 data_offset = data_length = data_sha256_hash = None 300 if data_blob is not None: 301 if do_hash_data_blob: 302 data_sha256_hash = hashlib.sha256(data_blob).digest() 303 data_length, data_offset = self.AddData(data_blob) 304 305 self.AddOperation(is_kernel, op_type, data_offset=data_offset, 306 data_length=data_length, src_extents=src_extents, 307 src_length=src_length, dst_extents=dst_extents, 308 dst_length=dst_length, data_sha256_hash=data_sha256_hash) 309 310 def WriteToFileWithData(self, file_obj, sigs_data=None, 311 privkey_file_name=None, 312 do_add_pseudo_operation=False, 313 is_pseudo_in_kernel=False, padding=None): 314 """Writes the payload content to a file, optionally signing the content. 315 316 Args: 317 file_obj: a file object open for writing 318 sigs_data: signatures blob to be appended to the payload (optional; 319 payload signature fields assumed to be preset by the caller) 320 privkey_file_name: key used for signing the payload (optional; used only 321 if explicit signatures blob not provided) 322 do_add_pseudo_operation: whether a pseudo-operation should be added to 323 account for the signature blob 324 is_pseudo_in_kernel: whether the pseudo-operation should be added to 325 kernel (True) or rootfs (False) operations 326 padding: stuff to dump past the normal data blobs provided (optional) 327 328 Raises: 329 TestError: if arguments are inconsistent or something goes wrong. 330 """ 331 sigs_len = len(sigs_data) if sigs_data else 0 332 333 # Do we need to generate a genuine signatures blob? 334 do_generate_sigs_data = sigs_data is None and privkey_file_name 335 336 if do_generate_sigs_data: 337 # First, sign some arbitrary data to obtain the size of a signature blob. 338 fake_sig = SignSha256('fake-payload-data', privkey_file_name) 339 fake_sigs_gen = SignaturesGenerator() 340 fake_sigs_gen.AddSig(1, fake_sig) 341 sigs_len = len(fake_sigs_gen.ToBinary()) 342 343 # Update the payload with proper signature attributes. 344 self.SetSignatures(self.curr_offset, sigs_len) 345 346 # Add a pseudo-operation to account for the signature blob, if requested. 347 if do_add_pseudo_operation: 348 if not self.block_size: 349 raise TestError('cannot add pseudo-operation without knowing the ' 350 'payload block size') 351 self.AddOperation( 352 is_pseudo_in_kernel, common.OpType.REPLACE, 353 data_offset=self.curr_offset, data_length=sigs_len, 354 dst_extents=[(common.PSEUDO_EXTENT_MARKER, 355 (sigs_len + self.block_size - 1) / self.block_size)]) 356 357 if do_generate_sigs_data: 358 # Once all payload fields are updated, dump and sign it. 359 temp_payload_file = cStringIO.StringIO() 360 self.WriteToFile(temp_payload_file, data_blobs=self.data_blobs) 361 sig = SignSha256(temp_payload_file.getvalue(), privkey_file_name) 362 sigs_gen = SignaturesGenerator() 363 sigs_gen.AddSig(1, sig) 364 sigs_data = sigs_gen.ToBinary() 365 assert len(sigs_data) == sigs_len, 'signature blob lengths mismatch' 366 367 # Dump the whole thing, complete with data and signature blob, to a file. 368 self.WriteToFile(file_obj, data_blobs=self.data_blobs, sigs_data=sigs_data, 369 padding=padding) 370