Adds support for optionally generating vbmeta.img in merge_builds.

Bug: 137853921
Bug: 138671115
Test: python -m unittest test_common
Test: python -m unittest test_add_img_to_target_files
Test: Ran 'merge_builds --build_vbmeta' for two devices, one with the
vbmeta struct on system.img and another with vbmeta_system.img. Flashed
the regenerated vbmeta.img files on devices, devices boot.

Change-Id: I8d7585c7af468be3d242d8aceeed6d27e6fc6d96
diff --git a/tools/releasetools/merge_builds.py b/tools/releasetools/merge_builds.py
index 7724d6f..ca348cf 100644
--- a/tools/releasetools/merge_builds.py
+++ b/tools/releasetools/merge_builds.py
@@ -24,9 +24,8 @@
 vendor partial build determines whether the merged result supports DAP.
 
 This script does not require builds to be built with 'make dist'.
-This script assumes that images other than super_empty.img do not require
-regeneration, including vbmeta images.
-TODO(b/137853921): Add support for regenerating vbmeta images.
+This script regenerates super_empty.img and vbmeta.img if necessary. Other
+images are assumed to not require regeneration.
 
 Usage: merge_builds.py [args]
 
@@ -39,6 +38,15 @@
 
   --product_out_vendor product_out_vendor_path
       Path to out/target/product/<vendor build>.
+
+  --build_vbmeta
+      If provided, vbmeta.img will be regenerated in out/target/product/<vendor
+      build>.
+
+  --framework_misc_info_keys
+      The optional path to a newline-separated config file containing keys to
+      obtain from the framework instance of misc_info.txt, used for creating
+      vbmeta.img. The remaining keys come from the vendor instance.
 """
 from __future__ import print_function
 
@@ -55,6 +63,8 @@
 OPTIONS.framework_images = ("system",)
 OPTIONS.product_out_framework = None
 OPTIONS.product_out_vendor = None
+OPTIONS.build_vbmeta = False
+OPTIONS.framework_misc_info_keys = None
 
 
 def CreateImageSymlinks():
@@ -82,6 +92,7 @@
   # super_empty.img from the framework build.
   if (framework_dict.get("use_dynamic_partitions") == "true") and (
       vendor_dict.get("use_dynamic_partitions") == "true"):
+    logger.info("Building super_empty.img.")
     merged_dict = dict(vendor_dict)
     merged_dict.update(
         common.MergeDynamicPartitionInfoDicts(
@@ -96,10 +107,52 @@
     build_super_image.BuildSuperImage(merged_dict, output_super_empty_path)
 
 
+def BuildVBMeta():
+  logger.info("Building vbmeta.img.")
+
+  framework_dict = common.LoadDictionaryFromFile(
+      os.path.join(OPTIONS.product_out_framework, "misc_info.txt"))
+  vendor_dict = common.LoadDictionaryFromFile(
+      os.path.join(OPTIONS.product_out_vendor, "misc_info.txt"))
+  merged_dict = dict(vendor_dict)
+  if OPTIONS.framework_misc_info_keys:
+    for key in common.LoadListFromFile(OPTIONS.framework_misc_info_keys):
+      merged_dict[key] = framework_dict[key]
+
+  # Build vbmeta.img using partitions in product_out_vendor.
+  partitions = {}
+  for partition in common.AVB_PARTITIONS:
+    partition_path = os.path.join(OPTIONS.product_out_vendor,
+                                  "%s.img" % partition)
+    if os.path.exists(partition_path):
+      partitions[partition] = partition_path
+
+  # vbmeta_partitions includes the partitions that should be included into
+  # top-level vbmeta.img, which are the ones that are not included in any
+  # chained VBMeta image plus the chained VBMeta images themselves.
+  vbmeta_partitions = common.AVB_PARTITIONS[:]
+  for partition in common.AVB_VBMETA_PARTITIONS:
+    chained_partitions = merged_dict.get("avb_%s" % partition, "").strip()
+    if chained_partitions:
+      partitions[partition] = os.path.join(OPTIONS.product_out_vendor,
+                                           "%s.img" % partition)
+      vbmeta_partitions = [
+          item for item in vbmeta_partitions
+          if item not in chained_partitions.split()
+      ]
+      vbmeta_partitions.append(partition)
+
+  output_vbmeta_path = os.path.join(OPTIONS.product_out_vendor, "vbmeta.img")
+  OPTIONS.info_dict = merged_dict
+  common.BuildVBMeta(output_vbmeta_path, partitions, "vbmeta",
+                     vbmeta_partitions)
+
+
 def MergeBuilds():
   CreateImageSymlinks()
   BuildSuperEmpty()
-  # TODO(b/137853921): Add support for regenerating vbmeta images.
+  if OPTIONS.build_vbmeta:
+    BuildVBMeta()
 
 
 def main():
@@ -112,6 +165,10 @@
       OPTIONS.product_out_framework = a
     elif o == "--product_out_vendor":
       OPTIONS.product_out_vendor = a
+    elif o == "--build_vbmeta":
+      OPTIONS.build_vbmeta = True
+    elif o == "--framework_misc_info_keys":
+      OPTIONS.framework_misc_info_keys = a
     else:
       return False
     return True
@@ -123,6 +180,8 @@
           "framework_images=",
           "product_out_framework=",
           "product_out_vendor=",
+          "build_vbmeta",
+          "framework_misc_info_keys=",
       ],
       extra_option_handler=option_handler)