Support signing custom images with AVB HASH or AVB HASHTREE

`make custom_images` supports to build different kinds of *non-droid* images,
e.g., odm.img. Adding the support of signing them with either AVB HASH footer
or AVB HASHTREE footer. The user can use HASH for small images and
HASHTREE for large images.

Sample signing configurations:
 * AVB HASH footer:
   - CUSTOM_IMAGE_AVB_HASH_ENABLE := true
   - CUSTOM_IMAGE_AVB_ADD_HASH_FOOTER_ARGS := --append_to_release_string my_odm_image

 * AVB HASHTREE footer:
   - CUSTOM_IMAGE_AVB_HASHTREE_ENABLE := true
   - CUSTOM_IMAGE_AVB_ADD_HASHTREE_FOOTER_ARGS := --fec_num_roots 8

 * Using custom signing key:
   - CUSTOM_IMAGE_AVB_ALGORITHM := SHA256_RSA2048
   - CUSTOM_IMAGE_AVB_KEY_PATH := external/avb/test/data/testkey_rsa2048.pem

Bug: 36701014
Test: `make custom_images` with AVB HASH footer
Test: `make custom_images` with AVB HASHTREE footer
Test: `make droid` to check system.img is still properly signed with AVB HASHTREE
Test: `make droid` to check vendor.img is still properly signed with AVB HASHTREE

Change-Id: I8dc420e12e37e9a631345c0cd883339db05d489f
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 3094dca..05c895a 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -97,18 +97,19 @@
   simg = sparse_img.SparseImage(image_file, mode="r+b", build_map=False)
   simg.AppendFillChunk(0, blocks)
 
-def AVBCalcMaxImageSize(avbtool, partition_size, additional_args):
+def AVBCalcMaxImageSize(avbtool, footer_type, partition_size, additional_args):
   """Calculates max image size for a given partition size.
 
   Args:
     avbtool: String with path to avbtool.
+    footer_type: 'hash' or 'hashtree' for generating footer.
     partition_size: The size of the partition in question.
     additional_args: Additional arguments to pass to 'avbtool
       add_hashtree_image'.
   Returns:
     The maximum image size or 0 if an error occurred.
   """
-  cmdline = "%s add_hashtree_footer " % avbtool
+  cmdline = "%s add_%s_footer " % (avbtool, footer_type)
   cmdline += "--partition_size %d " % partition_size
   cmdline += "--calc_max_image_size "
   cmdline += additional_args
@@ -118,13 +119,14 @@
   else:
     return int(output)
 
-def AVBAddHashtree(image_path, avbtool, partition_size, partition_name,
-                   signing_args, additional_args):
+def AVBAddFooter(image_path, avbtool, footer_type, partition_size,
+                 partition_name, signing_args, additional_args):
   """Adds dm-verity hashtree and AVB metadata to an image.
 
   Args:
     image_path: Path to image to modify.
     avbtool: String with path to avbtool.
+    footer_type: 'hash' or 'hashtree' for generating footer.
     partition_size: The size of the partition in question.
     partition_name: The name of the partition - will be embedded in metadata.
     signing_args: Arguments for signing the image.
@@ -133,7 +135,7 @@
   Returns:
     True if the operation succeeded.
   """
-  cmdline = "%s add_hashtree_footer " % avbtool
+  cmdline = "%s add_%s_footer " % (avbtool, footer_type)
   cmdline += "--partition_size %d " % partition_size
   cmdline += "--partition_name %s " % partition_name
   cmdline += "--image %s " % image_path
@@ -410,12 +412,19 @@
     prop_dict["original_partition_size"] = str(partition_size)
     prop_dict["verity_size"] = str(verity_size)
 
-  # Adjust partition size for AVB.
-  if prop_dict.get("avb_enable") == "true":
+  # Adjust partition size for AVB hash footer or AVB hashtree footer.
+  avb_footer_type = ''
+  if prop_dict.get("avb_hash_enable") == "true":
+    avb_footer_type = 'hash'
+  elif prop_dict.get("avb_hashtree_enable") == "true":
+    avb_footer_type = 'hashtree'
+
+  if avb_footer_type:
     avbtool = prop_dict.get("avb_avbtool")
     partition_size = int(prop_dict.get("partition_size"))
-    additional_args = prop_dict["avb_add_hashtree_footer_args"]
-    max_image_size = AVBCalcMaxImageSize(avbtool, partition_size,
+    # avb_add_hash_footer_args or avb_add_hashtree_footer_args.
+    additional_args = prop_dict["avb_add_" + avb_footer_type + "_footer_args"]
+    max_image_size = AVBCalcMaxImageSize(avbtool, avb_footer_type, partition_size,
                                          additional_args)
     if max_image_size == 0:
       return False
@@ -561,15 +570,16 @@
     if not MakeVerityEnabledImage(out_file, verity_fec_supported, prop_dict):
       return False
 
-  # Add AVB hashtree and metadata.
-  if "avb_enable" in prop_dict:
+  # Add AVB HASH or HASHTREE footer (metadata).
+  if avb_footer_type:
     avbtool = prop_dict.get("avb_avbtool")
     original_partition_size = int(prop_dict.get("original_partition_size"))
     partition_name = prop_dict["partition_name"]
     signing_args = prop_dict["avb_signing_args"]
-    additional_args = prop_dict["avb_add_hashtree_footer_args"]
-    if not AVBAddHashtree(out_file, avbtool, original_partition_size,
-                          partition_name, signing_args, additional_args):
+    # avb_add_hash_footer_args or avb_add_hashtree_footer_args
+    additional_args = prop_dict["avb_add_" + avb_footer_type + "_footer_args"]
+    if not AVBAddFooter(out_file, avbtool, avb_footer_type, original_partition_size,
+                      partition_name, signing_args, additional_args):
       return False
 
   if run_fsck and prop_dict.get("skip_fsck") != "true":
@@ -639,7 +649,7 @@
     copy_prop("system_squashfs_block_size", "squashfs_block_size")
     copy_prop("system_squashfs_disable_4k_align", "squashfs_disable_4k_align")
     copy_prop("system_base_fs_file", "base_fs_file")
-    copy_prop("system_avb_enable", "avb_enable")
+    copy_prop("system_avb_hashtree_enable", "avb_hashtree_enable")
     copy_prop("system_avb_add_hashtree_footer_args",
               "avb_add_hashtree_footer_args")
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
@@ -656,7 +666,7 @@
     copy_prop("system_squashfs_compressor_opt", "squashfs_compressor_opt")
     copy_prop("system_squashfs_block_size", "squashfs_block_size")
     copy_prop("system_base_fs_file", "base_fs_file")
-    copy_prop("system_avb_enable", "avb_enable")
+    copy_prop("system_avb_hashtree_enable", "avb_hashtree_enable")
     copy_prop("system_avb_add_hashtree_footer_args",
               "avb_add_hashtree_footer_args")
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
@@ -681,7 +691,7 @@
     copy_prop("vendor_squashfs_block_size", "squashfs_block_size")
     copy_prop("vendor_squashfs_disable_4k_align", "squashfs_disable_4k_align")
     copy_prop("vendor_base_fs_file", "base_fs_file")
-    copy_prop("vendor_avb_enable", "avb_enable")
+    copy_prop("vendor_avb_hashtree_enable", "avb_hashtree_enable")
     copy_prop("vendor_avb_add_hashtree_footer_args",
               "avb_add_hashtree_footer_args")
     copy_prop("vendor_extfs_inode_count", "extfs_inode_count")