diff --git a/avbroot.py b/avbroot.py index a93ca97..1906ed1 100644 --- a/avbroot.py +++ b/avbroot.py @@ -31,21 +31,30 @@ def print_status(*args, **kwargs): def get_images(manifest): boot_image = 'boot' - vendor_boot_image = None + otacert_image = None for p in manifest.partitions: # Devices launching with Android 13 use a GKI init_boot ramdisk if p.partition_name == 'init_boot': boot_image = p.partition_name + # OnePlus devices have a recovery image + elif p.partition_name == 'recovery': + # If a recovery image exists, it will always contain the OTA certs, + # even if vendor_boot exists + otacert_image = p.partition_name # Older devices may not have vendor_boot elif p.partition_name == 'vendor_boot': - vendor_boot_image = p.partition_name + if otacert_image is None: + otacert_image = p.partition_name images = ['vbmeta', boot_image] - if vendor_boot_image is not None: - images.append(vendor_boot_image) - return images, boot_image + if otacert_image is not None: + images.append(otacert_image) + else: + otacert_image = boot_image + + return images, boot_image, otacert_image def patch_ota_payload(f_in, f_out, file_size, magisk, privkey_avb, privkey_ota, @@ -59,22 +68,21 @@ def patch_ota_payload(f_in, f_out, file_size, magisk, privkey_avb, privkey_ota, os.mkdir(payload_dir) version, manifest, blob_offset = ota.parse_payload(f_in) - images, boot_image = get_images(manifest) + images, boot_image, otacert_image = get_images(manifest) print_status('Extracting', ', '.join(images), 'from the payload') ota.extract_images(f_in, manifest, blob_offset, extract_dir, images) boot_patches = [boot.MagiskRootPatch(magisk)] - vendor_boot_patches = [boot.OtaCertPatch(magisk, cert_ota)] + otacert_patches = [boot.OtaCertPatch(magisk, cert_ota)] - # Older devices don't have a vendor_boot - if 'vendor_boot' not in images: - boot_patches.extend(vendor_boot_patches) - vendor_boot_patches.clear() + if otacert_image == boot_image: + boot_patches.extend(otacert_patches) + otacert_patches.clear() avb = avbtool.Avb() - print_status('Patching boot image') + print_status(f'Patching {boot_image} image') boot.patch_boot( avb, os.path.join(extract_dir, f'{boot_image}.img'), @@ -84,15 +92,15 @@ def patch_ota_payload(f_in, f_out, file_size, magisk, privkey_avb, privkey_ota, boot_patches, ) - if vendor_boot_patches: - print_status('Patching vendor_boot image') + if otacert_patches: + print_status(f'Patching {otacert_image} image') boot.patch_boot( avb, - os.path.join(extract_dir, 'vendor_boot.img'), - os.path.join(patch_dir, 'vendor_boot.img'), + os.path.join(extract_dir, f'{otacert_image}.img'), + os.path.join(patch_dir, f'{otacert_image}.img'), privkey_avb, True, - vendor_boot_patches, + otacert_patches, ) print_status('Building new root vbmeta image') diff --git a/avbroot/boot.py b/avbroot/boot.py index b893ea9..d63443f 100644 --- a/avbroot/boot.py +++ b/avbroot/boot.py @@ -101,13 +101,20 @@ class OtaCertPatch(BootImagePatch): if magic == b'ANDROID!': f.seek(0x28) + is_vendor = False elif magic == b'VNDRBOOT': # Version is immediately after magic - pass + is_vendor = True else: raise Exception(b'Invalid magic: {magic}') - return struct.unpack('I', f.read(4))[0] + return struct.unpack('I', f.read(4))[0], is_vendor + + def _is_uncompressed(self, ramdisk_file): + with open(ramdisk_file, 'rb') as f: + magic = f.read(6) + + return magic == b'070701' or magic == b'070702' def patch(self, image_file, temp_dir): def run(*args): @@ -122,8 +129,10 @@ class OtaCertPatch(BootImagePatch): # ramdisks as there may be more than one. This is not the case for # Android 13 on the Pixel 6 Pro. ramdisk_file = 'ramdisk.cpio' - header_version = self._read_header_version(image_file) - if header_version == 4: + header_version, is_vendor = self._read_header_version(image_file) + need_decompress = header_version == 4 and is_vendor + + if need_decompress: ramdisk_file = 'decompressed.cpio' run('decompress', 'ramdisk.cpio', ramdisk_file) @@ -151,7 +160,7 @@ class OtaCertPatch(BootImagePatch): ) # Recompress ramdisk - if header_version == 4: + if need_decompress: run('compress=lz4_legacy', ramdisk_file, 'ramdisk.cpio') # Repack image diff --git a/avbroot/ota.py b/avbroot/ota.py index 420bce1..35e8ab7 100644 --- a/avbroot/ota.py +++ b/avbroot/ota.py @@ -92,7 +92,8 @@ def _extract_image(f_payload, f_out, block_size, blob_offset, partition): else: raise Exception(f'Unsupported operation: {op.type}') - if h_data.digest() != op.data_sha256_hash: + # OnePlus OTA payloads don't specify a checksum for ZERO blocks + if h_data.digest() != op.data_sha256_hash and op.type != Type.ZERO: raise Exception('Expected hash %s, but got %s' % (h_data.hexdigest(), binascii.hexlify(op.data_sha256_hash)))