Merge pull request #32 from chenxiaolong/oneplus

Add support for OnePlus OTAs

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')

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

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)))