Add options to sign the prebuilt custom images.

The custom images are any images owned by OEMs and SoCs, oem images
mounted on /oem is an example. The oem images can be used to customize
devices for different carriers, like wallpaper, ringtones, and
carrier-specific apks. OEMs can generate multiple oem images, like
oem.img, oem-carrier1.img and oem-carrier2.img and flash different oem
images for different carriers. The oem images are only one case, OEMs
and SoCs can add more custom images and mount them to custom partitions.

This change enables custom images to be vbmeta.img chained partitions.
The following configuration in BoardConfig.mk is an exmaple. It has two
custom partitions: oem and test. They will be signed by different keys.
And they will be chained by vbmeta.img. The custom images here are
prebuilts, which can be built by `make custom_images` separately.
BOARD_AVB_<CUSTOM_PARTITION>_IMAGE_LIST should include all custom images
to apply AVB signing. And to every custom partition, one image whose
name is partition name must be added in its
BOARD_AVB_<CUSTOM_PARTITION>_IMAGE_LIST.

BOARD_CUSTOMIMAGES_PARTITION_LIST := oem test

BOARD_AVB_OEM_KEY_PATH := external/avb/test/data/testkey_rsa4096.pem
BOARD_AVB_OEM_ALGORITHM := SHA256_RSA4096
BOARD_AVB_OEM_ADD_HASHTREE_FOOTER_ARGS :=
BOARD_AVB_OEM_ROLLBACK_INDEX_LOCATION := 1
BOARD_AVB_OEM_PARTITION_SIZE := 5242880
BOARD_AVB_OEM_IMAGE_LIST := \
    device/xxxx/yyyy/oem/oem.img \
    device/xxxx/yyyy/oem/oem1.img

BOARD_AVB_TEST_KEY_PATH := external/avb/test/data/testkey_rsa2048.pem
BOARD_AVB_TEST_ALGORITHM := SHA256_RSA2048
BOARD_AVB_TEST_ADD_HASHTREE_FOOTER_ARGS :=
BOARD_AVB_TEST_ROLLBACK_INDEX_LOCATION := 2
BOARD_AVB_TEST_PARTITION_SIZE := 10485760
BOARD_AVB_TEST_IMAGE_LIST := \
    device/xxxx/yyyy/test/test.img \
    device/xxxx/yyyy/test/test1.img

To resign the custom images in the target zip file, the
avb_extra_custom_image_key, avb_extra_custom_image_algorithms and
avb_extra_custom_image_extra_args options are added to the
sign_target_files_apks tool too. The following test cases list some
examples about how to use them.

BUG: 154171021
Test: 1) "atest --host releasetools_test releasetools_py3_test -c"
  2) Build images by 'make dist', sign and validate target files.
   a) Test on dist w/ chained vbmeta_system and ome custom images
   sign_target_files_apks -d certs \
     --avb_extra_custom_image_key oem=oem_rsa4096.pem \
     --avb_extra_custom_image_algorithm oem=SHA256_RSA4096 \
     xxx-target_xxx.zip signed.zip
   validate_target_files.py signed.zip
   Flash image and boot up.
   Verify the oem images and vbmeta images in OUT and target zips by
   avbtool.

   b) Test on dist w/ chained vbmeta_system and oem and test custom images
   sign_target_files_apks -d certs \
     --avb_extra_custom_image_key oem=oem_rsa4096.pem \
     --avb_extra_custom_image_algorithm oem=SHA256_RSA4096 \
     --avb_extra_custom_image_extra_args oem=--do_not_generate_fec \
     --avb_extra_custom_image_key test=test_rsa4096.pem \
     --avb_extra_custom_image_algorithm test=SHA256_RSA4096 \
     xxx-target_xxx.zip signed.zip
   validate_target_files.py signed.zip
   Verify the oem, test images and vbmeta images in OUT and target zips
   by avbtool.

   c) Test on dist w/o chained partition.
   sign_target_files_apks -d certs xxx-target_xxx.zip signed.zip
   validate_target_files.py signed.zip
   Flash image and boot up.
   Verify the vbmeta images in OUT and target zips by avbtool.

Change-Id: Ifccfee5e8909697eef6ccda0cc352fa16a9f6db6
diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index 52cd9a8..52b7889 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -112,6 +112,17 @@
       (e.g. "--signing_helper /path/to/helper"). The args will be appended to
       the existing ones in info dict.
 
+  --avb_extra_custom_image_key <partition=key>
+  --avb_extra_custom_image_algorithm <partition=algorithm>
+      Use the specified algorithm (e.g. SHA256_RSA4096) and the key to AVB-sign
+      the specified custom images mounted on the partition. Otherwise it uses
+      the existing values in info dict.
+
+  --avb_extra_custom_image_extra_args <partition=extra_args>
+      Specify any additional args that are needed to AVB-sign the custom images
+      mounted on the partition (e.g. "--signing_helper /path/to/helper"). The
+      args will be appended to the existing ones in info dict.
+
   --android_jar_path <path>
       Path to the android.jar to repack the apex file.
 """
@@ -956,12 +967,20 @@
     if extra_args:
       print('Setting extra AVB signing args for %s to "%s"' % (
           partition, extra_args))
-      args_key = AVB_FOOTER_ARGS_BY_PARTITION[partition]
+      if partition in AVB_FOOTER_ARGS_BY_PARTITION:
+        args_key = AVB_FOOTER_ARGS_BY_PARTITION[partition]
+      else:
+        # custom partition
+        args_key = "avb_{}_add_hashtree_footer_args".format(partition)
       misc_info[args_key] = (misc_info.get(args_key, '') + ' ' + extra_args)
 
   for partition in AVB_FOOTER_ARGS_BY_PARTITION:
     ReplaceAvbPartitionSigningKey(partition)
 
+  for custom_partition in misc_info.get(
+      "avb_custom_images_partition_list", "").strip().split():
+    ReplaceAvbPartitionSigningKey(custom_partition)
+
 
 def RewriteAvbProps(misc_info):
   """Rewrites the props in AVB signing args."""
@@ -1208,6 +1227,15 @@
       OPTIONS.avb_extra_args['vbmeta_vendor'] = a
     elif o == "--avb_apex_extra_args":
       OPTIONS.avb_extra_args['apex'] = a
+    elif o == "--avb_extra_custom_image_key":
+      partition, key = a.split("=")
+      OPTIONS.avb_keys[partition] = key
+    elif o == "--avb_extra_custom_image_algorithm":
+      partition, algorithm = a.split("=")
+      OPTIONS.avb_algorithms[partition] = algorithm
+    elif o == "--avb_extra_custom_image_extra_args":
+      partition, extra_args = a.split("=")
+      OPTIONS.avb_extra_args[partition] = extra_args
     else:
       return False
     return True
@@ -1252,6 +1280,9 @@
           "avb_vbmeta_vendor_algorithm=",
           "avb_vbmeta_vendor_key=",
           "avb_vbmeta_vendor_extra_args=",
+          "avb_extra_custom_image_key=",
+          "avb_extra_custom_image_algorithm=",
+          "avb_extra_custom_image_extra_args=",
       ],
       extra_option_handler=option_handler)