OTA: Support A/B devices custom images update.

Add a new custom_image option to configure which custom images to use to
update custom partitions in A/B update.

This change also moves oem_settings to common option as A/B update will
use it to set oem properties too.

BUG: 171225290
Test: unittest pass, generate OTAs, flash to devices and check results
Change-Id: I279477d6b2954fb3705d7efede0a8bcd330c108b

tools/releasetools/ota_from_target_files.py

@@ -85,6 +85,13 @@ Common options that apply to both of non-A/B and A/B OTAs
       If not set, generates A/B package for A/B device and non-A/B package for
       non-A/B device.
 
+  -o  (--oem_settings) <main_file[,additional_files...]>
+      Comma separated list of files used to specify the expected OEM-specific
+      properties on the OEM partition of the intended device. Multiple expected
+      values can be used by providing multiple files. Only the first dict will
+      be used to compute fingerprint, while the rest will be used to assert
+      OEM-specific properties.
+
 Non-A/B OTA specific options
 
   -b  (--binary) <file>
@@ -114,13 +121,6 @@ Non-A/B OTA specific options
       builds for an incremental package. This option is only meaningful when -i
       is specified.
 
-  -o  (--oem_settings) <main_file[,additional_files...]>
-      Comma seperated list of files used to specify the expected OEM-specific
-      properties on the OEM partition of the intended device. Multiple expected
-      values can be used by providing multiple files. Only the first dict will
-      be used to compute fingerprint, while the rest will be used to assert
-      OEM-specific properties.
-
   --oem_no_mount
       For devices with OEM-specific properties but without an OEM partition, do
       not mount the OEM partition in the updater-script. This should be very
@@ -206,6 +206,11 @@ A/B OTA specific options
   --partial "<PARTITION> [<PARTITION>[...]]"
       Generate partial updates, overriding ab_partitions list with the given
       list.
+
+  --custom_image <custom_partition=custom_image>
+      Use the specified custom_image to update custom_partition when generating
+      an A/B OTA package. e.g. "--custom_image oem=oem.img --custom_image
+      cus=cus_test.img"
 """
 
 from __future__ import print_function
@@ -262,7 +267,7 @@ OPTIONS.skip_postinstall = False
 OPTIONS.skip_compatibility_check = False
 OPTIONS.disable_fec_computation = False
 OPTIONS.partial = None
-
+OPTIONS.custom_images = {}
 
 POSTINSTALL_CONFIG = 'META/postinstall_config.txt'
 DYNAMIC_PARTITION_INFO = 'META/dynamic_partitions_info.txt'
@@ -901,6 +906,43 @@ def GetTargetFilesZipForRetrofitDynamicPartitions(input_file,
 
   return target_file
 
+def GetTargetFilesZipForCustomImagesUpdates(input_file, custom_images):
+  """Returns a target-files.zip for custom partitions update.
+
+  This function modifies ab_partitions list with the desired custom partitions
+  and puts the custom images into the target target-files.zip.
+
+  Args:
+    input_file: The input target-files.zip filename.
+    custom_images: A map of custom partitions and custom images.
+
+  Returns:
+    The filename of a target-files.zip which has renamed the custom images in
+    the IMAGS/ to their partition names.
+  """
+  # Use zip2zip to avoid extracting the zipfile.
+  target_file = common.MakeTempFile(prefix="targetfiles-", suffix=".zip")
+  cmd = ['zip2zip', '-i', input_file, '-o', target_file]
+
+  with zipfile.ZipFile(input_file, allowZip64=True) as input_zip:
+    namelist = input_zip.namelist()
+
+  # Write {custom_image}.img as {custom_partition}.img.
+  for custom_partition, custom_image in custom_images.items():
+    default_custom_image = '{}.img'.format(custom_partition)
+    if default_custom_image != custom_image:
+      logger.info("Update custom partition '%s' with '%s'",
+                  custom_partition, custom_image)
+      # Default custom image need to be deleted first.
+      namelist.remove('IMAGES/{}'.format(default_custom_image))
+      # IMAGES/{custom_image}.img:IMAGES/{custom_partition}.img.
+      cmd.extend(['IMAGES/{}:IMAGES/{}'.format(custom_image,
+                                               default_custom_image)])
+
+  cmd.extend(['{}:{}'.format(name, name) for name in namelist])
+  common.RunAndCheckOutput(cmd)
+
+  return target_file
 
 def GenerateAbOtaPackage(target_file, output_file, source_file=None):
   """Generates an Android OTA package that has A/B update payload."""
@@ -927,6 +969,11 @@ def GenerateAbOtaPackage(target_file, output_file, source_file=None):
 
   additional_args = []
 
+  # Prepare custom images.
+  if OPTIONS.custom_images:
+    target_file = GetTargetFilesZipForCustomImagesUpdates(
+        target_file, OPTIONS.custom_images)
+
   if OPTIONS.retrofit_dynamic_partitions:
     target_file = GetTargetFilesZipForRetrofitDynamicPartitions(
         target_file, target_info.get("super_block_devices").strip().split(),
@@ -1105,6 +1152,9 @@ def main(argv):
       if not partitions:
         raise ValueError("Cannot parse partitions in {}".format(a))
       OPTIONS.partial = partitions
+    elif o == "--custom_image":
+      custom_partition, custom_image = a.split("=")
+      OPTIONS.custom_images[custom_partition] = custom_image
     else:
       return False
     return True
@@ -1144,6 +1194,7 @@ def main(argv):
                                  "force_non_ab",
                                  "boot_variable_file=",
                                  "partial=",
+                                 "custom_image=",
                              ], extra_option_handler=option_handler)
 
   if len(args) != 2:

tools/releasetools/test_ota_from_target_files.py

@@ -27,6 +27,7 @@ from ota_utils import (
     FinalizeMetadata, GetPackageMetadata, PropertyFiles)
 from ota_from_target_files import (
     _LoadOemDicts, AbOtaPropertyFiles,
+    GetTargetFilesZipForCustomImagesUpdates,
     GetTargetFilesZipForPartialUpdates,
     GetTargetFilesZipForSecondaryImages,
     GetTargetFilesZipWithoutPostinstallConfig,
@@ -545,6 +546,46 @@ class OtaFromTargetFilesTest(test_utils.ReleaseToolsTestCase):
     with zipfile.ZipFile(target_file) as verify_zip:
       self.assertNotIn(POSTINSTALL_CONFIG, verify_zip.namelist())
 
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_GetTargetFilesZipForCustomImagesUpdates_oemDefaultImage(self):
+    input_file = construct_target_files()
+    with zipfile.ZipFile(input_file, 'a', allowZip64=True) as append_zip:
+      common.ZipWriteStr(append_zip, 'IMAGES/oem.img', 'oem')
+      common.ZipWriteStr(append_zip, 'IMAGES/oem_test.img', 'oem_test')
+
+    target_file = GetTargetFilesZipForCustomImagesUpdates(
+        input_file, {'oem': 'oem.img'})
+
+    with zipfile.ZipFile(target_file) as verify_zip:
+      namelist = verify_zip.namelist()
+      ab_partitions = verify_zip.read('META/ab_partitions.txt').decode()
+      oem_image = verify_zip.read('IMAGES/oem.img').decode()
+
+    self.assertIn('META/ab_partitions.txt', namelist)
+    self.assertEqual('boot\nsystem\nvendor\nbootloader\nmodem', ab_partitions)
+    self.assertIn('IMAGES/oem.img', namelist)
+    self.assertEqual('oem', oem_image)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_GetTargetFilesZipForCustomImagesUpdates_oemTestImage(self):
+    input_file = construct_target_files()
+    with zipfile.ZipFile(input_file, 'a', allowZip64=True) as append_zip:
+      common.ZipWriteStr(append_zip, 'IMAGES/oem.img', 'oem')
+      common.ZipWriteStr(append_zip, 'IMAGES/oem_test.img', 'oem_test')
+
+    target_file = GetTargetFilesZipForCustomImagesUpdates(
+        input_file, {'oem': 'oem_test.img'})
+
+    with zipfile.ZipFile(target_file) as verify_zip:
+      namelist = verify_zip.namelist()
+      ab_partitions = verify_zip.read('META/ab_partitions.txt').decode()
+      oem_image = verify_zip.read('IMAGES/oem.img').decode()
+
+    self.assertIn('META/ab_partitions.txt', namelist)
+    self.assertEqual('boot\nsystem\nvendor\nbootloader\nmodem', ab_partitions)
+    self.assertIn('IMAGES/oem.img', namelist)
+    self.assertEqual('oem_test', oem_image)
+
   def _test_FinalizeMetadata(self, large_entry=False):
     entries = [
         'required-entry1',