diff --git a/avbroot.py b/avbroot.py index 9d8a00e..3845f70 100644 --- a/avbroot.py +++ b/avbroot.py @@ -18,14 +18,8 @@ from avbroot import util from avbroot import vbmeta -PATH_METADATA_PB = 'META-INF/com/android/metadata.pb' PATH_PAYLOAD = 'payload.bin' PATH_PROPERTIES = 'payload_properties.txt' -SKIP_PATHS = ( - PATH_METADATA_PB, - 'META-INF/com/android/metadata', - 'META-INF/com/android/otacert', -) def print_status(*args, **kwargs): @@ -130,7 +124,7 @@ def patch_ota_zip(f_zip_in, f_zip_out, magisk, privkey_avb, privkey_ota, zipfile.ZipFile(f_zip_out, 'w') as z_out, ): infolist = z_in.infolist() - missing = {PATH_METADATA_PB, PATH_PAYLOAD, PATH_PROPERTIES} + missing = {ota.PATH_METADATA_PB, PATH_PAYLOAD, PATH_PROPERTIES} i_payload = -1 i_properties = -1 @@ -159,15 +153,10 @@ def patch_ota_zip(f_zip_in, f_zip_out, magisk, privkey_avb, privkey_ota, for info in z_in.infolist(): # The existing metadata is needed to generate a new signed zip - if info.filename == PATH_METADATA_PB: + if info.filename == ota.PATH_METADATA_PB: with z_in.open(info, 'r') as f_in: metadata = f_in.read() - # Skip files that are created during zip signing - if info.filename in SKIP_PATHS: - print_status('Skipping', info.filename) - continue - # Copy other files, patching if needed with ( z_in.open(info, 'r') as f_in, @@ -222,39 +211,34 @@ def patch_subcommand(args): dec_privkey_avb = os.path.join(key_dir, 'avb.key') openssl.decrypt_key(args.privkey_avb, dec_privkey_avb) - # AOSP's OTA utils require a DER-encoded private key with the `.pk8` - # extension and a PEM-encoded certificate with the `.x509.pem` extension - dec_privkey_ota = os.path.join(key_dir, 'ota.pk8') - key_prefix_ota = dec_privkey_ota[:-4] + # signapk requires a DER-encoded private key + dec_privkey_ota = os.path.join(key_dir, 'ota.key') openssl.decrypt_key(args.privkey_ota, dec_privkey_ota, out_form='DER') - cert_ota = os.path.join(key_dir, 'ota.x509.pem') - shutil.copyfile(args.cert_ota, cert_ota) - # Ensure that the certificate matches the private key - if not openssl.cert_matches_key(cert_ota, dec_privkey_ota): + if not openssl.cert_matches_key(args.cert_ota, dec_privkey_ota): raise Exception('OTA certificate does not match private key') start = time.perf_counter_ns() - with tempfile.NamedTemporaryFile() as temp_unsigned: + with util.open_output_file(output) as temp: metadata = patch_ota_zip( args.input, - temp_unsigned, + temp, args.magisk, dec_privkey_avb, dec_privkey_ota, - cert_ota, + args.cert_ota, ) print_status('Signing OTA zip') - with util.open_output_file(output) as temp_signed: - ota.sign_zip( - temp_unsigned.name, - temp_signed.name, - key_prefix_ota, - metadata, - ) + ota.sign_zip( + temp.name, + temp.name, + dec_privkey_ota, + args.cert_ota, + metadata, + ) # Excluding the time it takes for the user to type in the passwords elapsed = time.perf_counter_ns() - start diff --git a/avbroot/ota.py b/avbroot/ota.py index 421e70f..bfbd172 100644 --- a/avbroot/ota.py +++ b/avbroot/ota.py @@ -2,19 +2,22 @@ import base64 import binascii import bz2 import hashlib +import itertools import lzma import os +import shutil import struct import sys +import subprocess +import zipfile # Silence undesired warning orig_argv0 = sys.argv[0] sys.argv[0] = sys.argv[0].removesuffix('.py') -import common +import ota_utils sys.argv[0] = orig_argv0 import ota_metadata_pb2 -import ota_utils import update_metadata_pb2 from . import openssl @@ -23,6 +26,9 @@ from . import util OTA_MAGIC = b'CrAU' +PATH_METADATA = 'META-INF/com/android/metadata' +PATH_METADATA_PB = f'{PATH_METADATA}.pb' + def parse_payload(f): ''' @@ -77,12 +83,12 @@ def _extract_image(f_payload, f_out, block_size, blob_offset, partition): hasher=h_data) elif op.type == Type.REPLACE_BZ: decompressor = bz2.BZ2Decompressor() - util.decompress_n(decompressor, f_payload, f_out, op.data_length, - hasher=h_data) + util.decompress_n(decompressor, f_payload, f_out, + op.data_length, hasher=h_data) elif op.type == Type.REPLACE_XZ: decompressor = lzma.LZMADecompressor() - util.decompress_n(decompressor, f_payload, f_out, op.data_length, - hasher=h_data) + util.decompress_n(decompressor, f_payload, f_out, + op.data_length, hasher=h_data) elif op.type == Type.ZERO or op.type == Type.DISCARD: util.zero_n(f_out, extent.num_blocks * block_size) else: @@ -258,7 +264,7 @@ def patch_payload(f_in, f_out, version, manifest, blob_offset, temp_dir, for name, path in patched.items(): # Find the partition in the manifest partition = next((p for p in manifest.partitions - if p.partition_name == name), None) + if p.partition_name == name), None) if partition is None: raise Exception(f'Partition {name} not found in manifest') @@ -341,7 +347,7 @@ def patch_payload(f_in, f_out, version, manifest, blob_offset, temp_dir, blob_size = manifest.signatures_offset + manifest.signatures_size new_file_size = metadata_size + dummy_sig_size + blob_size - b64 = lambda d: base64.b64encode(d) + def b64(d): return base64.b64encode(d) props = [ b'FILE_HASH=%s\n' % b64(h_full.digest()), b'FILE_SIZE=%d\n' % new_file_size, @@ -352,12 +358,62 @@ def patch_payload(f_in, f_out, version, manifest, blob_offset, temp_dir, return b''.join(props) -def sign_zip(input_path, output_path, key_prefix, metadata_raw): +def _replace_metadata(z_in, z_out, metadata): + ''' + Copy all entries from the input to the output, replacing the existing + metadata files (legacy and protobuf) with the data serialized from the + given metadata instance. + ''' + + legacy_metadata = ota_utils.BuildLegacyOtaMetadata(metadata) + legacy_metadata_str = "".join([f'{k}={v}\n' for k, v in + sorted(legacy_metadata.items())]) + metadata_bytes = metadata.SerializeToString() + + for info in z_in.infolist(): + with ( + z_in.open(info, 'r') as f_in, + z_out.open(info, 'w') as f_out, + ): + if info.filename == PATH_METADATA: + f_out.write(legacy_metadata_str.encode('UTF-8')) + elif info.filename == PATH_METADATA_PB: + f_out.write(metadata_bytes) + else: + shutil.copyfileobj(f_in, f_out) + + +def _sign_file(input_path, output_path, privkey, cert): + ''' + Sign an OTA zip with AOSP's signapk tool. This may result in the zip + entries being reordered. + ''' + + subprocess.check_call([ + 'java', + '-Xmx4096m', + '-jar', + os.path.join( + os.path.dirname(__file__), + '..', 'signapk', 'build', 'libs', 'signapk-all.jar', + ), + '-w', + cert, + privkey, + input_path, + output_path, + ]) + + +def sign_zip(input_path, output_path, privkey_ota, cert_ota, metadata_raw): ''' Sign an unsigned OTA zip. should be the serialized OTA - metadata protobuf struct from the original OTA. The property files contained - within the metadata that reference stored zip entries will be deleted and - recreated during signing. + metadata protobuf struct from the original OTA. The property files + contained within the metadata that reference stored zip entries will be + deleted and recreated during signing. + + This is a reimplementation of ota_utils.FinalizeMetadata() because that + function is quite inefficient, disk space and memory wise. ''' metadata = ota_metadata_pb2.OtaMetadata() @@ -365,32 +421,67 @@ def sign_zip(input_path, output_path, key_prefix, metadata_raw): metadata.property_files.clear() - # We can't replace common.OPTIONS itself because ota_utils holds a reference - # to the original instance - attrs = ( - 'search_path', - 'signapk_shared_library_path', - 'signapk_path', - ) - orig_attrs = {a: getattr(common.OPTIONS, a) for a in attrs} + props = [ + ota_utils.AbOtaPropertyFiles(), + ota_utils.StreamingPropertyFiles(), + ] - try: - common.OPTIONS.search_path = '/var/empty' - common.OPTIONS.signapk_shared_library_path = '/var/empty' - common.OPTIONS.signapk_path = os.path.join( - os.path.dirname(__file__), - '..', 'signapk', 'build', 'libs', 'signapk-all.jar', - ) + # The first signing attempt uses the specified input file initially and + # writes to the specified output file. If more signing attempts are needed, + # the now-intermediate output file will be used as the input for the next + # attempt. + attempt_input_path = input_path - ota_utils.FinalizeMetadata( - metadata, - input_path, - output_path, - package_key=key_prefix, - ) + # This is set to true once we know that the metadata entries can fit in the + # reserved space within the property files. + reserved_ok = False - finally: - common.Cleanup() + for attempt in itertools.count(): + # Compute initial property files with reserved space as placeholders + # to store the metadata entries later + if not reserved_ok: + with zipfile.ZipFile(attempt_input_path, 'r') as z: + for p in props: + metadata.property_files[p.name] = p.Compute(z) - for k, v in orig_attrs.items(): - setattr(common.OPTIONS, k, v) + # Replace the metadata files + with ( + util.open_output_file(output_path) as f_replace, + zipfile.ZipFile(attempt_input_path, 'r') as z_in, + zipfile.ZipFile(f_replace, 'w') as z_out, + ): + _replace_metadata(z_in, z_out, metadata) + + # The newly generated zip is now the input for further attempts + attempt_input_path = output_path + + # Sign the zip "in place" + with util.open_output_file(output_path) as f_replace: + _sign_file(output_path, f_replace.name, privkey_ota, cert_ota) + + with zipfile.ZipFile(output_path, 'r') as z: + # Compute the final property files, including metadata entries + try: + for p in props: + metadata.property_files[p.name] = \ + p.Finalize(z, len(metadata.property_files[p.name])) + + reserved_ok = True + except ota_utils.PropertyFiles.InsufficientSpaceException: + # There is not enough space in the placeholders to store the + # additional metadata entries and the new file offsets after + # potential reordering by signapk. Try again since reordering + # should no longer happen. + if attempt == 0: + continue + else: + raise + + # Attempt number 0 only has placeholders, so don't bother verifying + if attempt > 0: + # Verify that the zip entry offsets in the property files are + # still valid after signing + for p in props: + p.Verify(z, metadata.property_files[p.name].strip()) + + break