Build vendor images first when rebuilding with vendor_otatools.

Only supports rebuilding vendor.img and odm.img currently.

This allows "other" images (e.g. vbmeta, super_empty, ...) to
be rebuilt using the latest tooling.

Bug: 188491126
Bug: 192422274
Test: Sign a GRF R+S merged target_files package using:
      sign_target_files_apks \
        --vendor_otatools=<otatools.zip from R> \
	--vendor_partitions=vendor,odm \
	merged-target_files.zip \
	signed-target_files.zip
Change-Id: Ib93aea9f083ee9b722c31a42dcc780b5222053b8
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index f3b58f8..01cc233 100644
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -589,7 +589,7 @@
     AssertionError: If it can't find an image.
   """
   for partition in ab_partitions:
-    img_name = partition.strip() + ".img"
+    img_name = partition + ".img"
 
     # Assert that the image is present under IMAGES/ now.
     if output_zip:
@@ -856,39 +856,23 @@
         if output_zip:
           recovery_two_step_image.AddToZip(output_zip)
 
-  if has_system:
-    banner("system")
-    partitions['system'] = AddSystem(
-        output_zip, recovery_img=recovery_image, boot_img=boot_image)
+  def add_partition(partition, has_partition, add_func, add_args):
+    if has_partition:
+      banner(partition)
+      partitions[partition] = add_func(output_zip, *add_args)
 
-  if has_vendor:
-    banner("vendor")
-    partitions['vendor'] = AddVendor(
-        output_zip, recovery_img=recovery_image, boot_img=boot_image)
-
-  if has_product:
-    banner("product")
-    partitions['product'] = AddProduct(output_zip)
-
-  if has_system_ext:
-    banner("system_ext")
-    partitions['system_ext'] = AddSystemExt(output_zip)
-
-  if has_odm:
-    banner("odm")
-    partitions['odm'] = AddOdm(output_zip)
-
-  if has_vendor_dlkm:
-    banner("vendor_dlkm")
-    partitions['vendor_dlkm'] = AddVendorDlkm(output_zip)
-
-  if has_odm_dlkm:
-    banner("odm_dlkm")
-    partitions['odm_dlkm'] = AddOdmDlkm(output_zip)
-
-  if has_system_other:
-    banner("system_other")
-    AddSystemOther(output_zip)
+  add_partition_calls = (
+      ("system", has_system, AddSystem, [recovery_image, boot_image]),
+      ("vendor", has_vendor, AddVendor, [recovery_image, boot_image]),
+      ("product", has_product, AddProduct, []),
+      ("system_ext", has_system_ext, AddSystemExt, []),
+      ("odm", has_odm, AddOdm, []),
+      ("vendor_dlkm", has_vendor_dlkm, AddVendorDlkm, []),
+      ("odm_dlkm", has_odm_dlkm, AddOdmDlkm, []),
+      ("system_other", has_system_other, AddSystemOther, []),
+  )
+  for call in add_partition_calls:
+    add_partition(*call)
 
   AddApexInfo(output_zip)
 
@@ -902,13 +886,10 @@
     banner("partition-table")
     AddPartitionTable(output_zip)
 
-  if OPTIONS.info_dict.get("has_dtbo") == "true":
-    banner("dtbo")
-    partitions['dtbo'] = AddDtbo(output_zip)
-
-  if OPTIONS.info_dict.get("has_pvmfw") == "true":
-    banner("pvmfw")
-    partitions['pvmfw'] = AddPvmfw(output_zip)
+  add_partition("dtbo",
+                OPTIONS.info_dict.get("has_dtbo") == "true", AddDtbo, [])
+  add_partition("pvmfw",
+                OPTIONS.info_dict.get("has_pvmfw") == "true", AddPvmfw, [])
 
   # Custom images.
   custom_partitions = OPTIONS.info_dict.get(
@@ -965,7 +946,7 @@
                                    "ab_partitions.txt")
   if os.path.exists(ab_partitions_txt):
     with open(ab_partitions_txt) as f:
-      ab_partitions = f.readlines()
+      ab_partitions = f.read().splitlines()
 
     # For devices using A/B update, make sure we have all the needed images
     # ready under IMAGES/ or RADIO/.
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index c6800e8..5affa32 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -1925,14 +1925,14 @@
   RunAndCheckOutput(cmd)
 
 
-def UnzipTemp(filename, pattern=None):
+def UnzipTemp(filename, patterns=None):
   """Unzips the given archive into a temporary directory and returns the name.
 
   Args:
     filename: If filename is of the form "foo.zip+bar.zip", unzip foo.zip into
     a temp dir, then unzip bar.zip into that_dir/BOOTABLE_IMAGES.
 
-    pattern: Files to unzip from the archive. If omitted, will unzip the entire
+    patterns: Files to unzip from the archive. If omitted, will unzip the entire
     archvie.
 
   Returns:
@@ -1942,11 +1942,11 @@
   tmp = MakeTempDir(prefix="targetfiles-")
   m = re.match(r"^(.*[.]zip)\+(.*[.]zip)$", filename, re.IGNORECASE)
   if m:
-    UnzipToDir(m.group(1), tmp, pattern)
-    UnzipToDir(m.group(2), os.path.join(tmp, "BOOTABLE_IMAGES"), pattern)
+    UnzipToDir(m.group(1), tmp, patterns)
+    UnzipToDir(m.group(2), os.path.join(tmp, "BOOTABLE_IMAGES"), patterns)
     filename = m.group(1)
   else:
-    UnzipToDir(filename, tmp, pattern)
+    UnzipToDir(filename, tmp, patterns)
 
   return tmp
 
diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index 0842af9..d219ed6 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -189,6 +189,8 @@
 OPTIONS.gki_signing_algorithm = None
 OPTIONS.gki_signing_extra_args = None
 OPTIONS.android_jar_path = None
+OPTIONS.vendor_partitions = set()
+OPTIONS.vendor_otatools = None
 
 
 AVB_FOOTER_ARGS_BY_PARTITION = {
@@ -216,6 +218,10 @@
   if partition not in AVB_FOOTER_ARGS_BY_PARTITION:
     raise RuntimeError("Missing {} in AVB_FOOTER_ARGS".format(partition))
 
+# Partitions that can be regenerated after signing using a separate
+# vendor otatools package.
+ALLOWED_VENDOR_PARTITIONS = set(["vendor", "odm"])
+
 
 def IsApexFile(filename):
   return filename.endswith(".apex") or filename.endswith(".capex")
@@ -1180,6 +1186,63 @@
   return keys
 
 
+def BuildVendorPartitions(output_zip_path):
+  """Builds OPTIONS.vendor_partitions using OPTIONS.vendor_otatools."""
+  if OPTIONS.vendor_partitions.difference(ALLOWED_VENDOR_PARTITIONS):
+    logger.warning("Allowed --vendor_partitions: %s",
+                   ",".join(ALLOWED_VENDOR_PARTITIONS))
+    OPTIONS.vendor_partitions = ALLOWED_VENDOR_PARTITIONS.intersection(
+        OPTIONS.vendor_partitions)
+
+  logger.info("Building vendor partitions using vendor otatools.")
+  vendor_tempdir = common.UnzipTemp(output_zip_path, [
+      "META/*",
+  ] + ["{}/*".format(p.upper()) for p in OPTIONS.vendor_partitions])
+
+  # Disable various partitions that build based on misc_info fields.
+  # Only partitions in ALLOWED_VENDOR_PARTITIONS can be rebuilt using
+  # vendor otatools. These other partitions will be rebuilt using the main
+  # otatools if necessary.
+  vendor_misc_info_path = os.path.join(vendor_tempdir, "META/misc_info.txt")
+  vendor_misc_info = common.LoadDictionaryFromFile(vendor_misc_info_path)
+  vendor_misc_info["no_boot"] = "true"  # boot
+  vendor_misc_info["vendor_boot"] = "false"  # vendor_boot
+  vendor_misc_info["no_recovery"] = "true"  # recovery
+  vendor_misc_info["board_bpt_enable"] = "false"  # partition-table
+  vendor_misc_info["has_dtbo"] = "false"  # dtbo
+  vendor_misc_info["has_pvmfw"] = "false"  # pvmfw
+  vendor_misc_info["avb_custom_images_partition_list"] = ""  # custom images
+  vendor_misc_info["avb_enable"] = "false"  # vbmeta
+  vendor_misc_info["use_dynamic_partitions"] = "false"  # super_empty
+  vendor_misc_info["build_super_partition"] = "false"  # super split
+  with open(vendor_misc_info_path, "w") as output:
+    for key in sorted(vendor_misc_info):
+      output.write("{}={}\n".format(key, vendor_misc_info[key]))
+
+  # Disable care_map.pb as not all ab_partitions are available when
+  # vendor otatools regenerates vendor images.
+  os.remove(os.path.join(vendor_tempdir, "META/ab_partitions.txt"))
+
+  # Build vendor images using vendor otatools.
+  vendor_otatools_dir = common.MakeTempDir(prefix="vendor_otatools_")
+  common.UnzipToDir(OPTIONS.vendor_otatools, vendor_otatools_dir)
+  cmd = [
+      os.path.join(vendor_otatools_dir, "bin", "add_img_to_target_files"),
+      "--is_signing",
+      "--verbose",
+      vendor_tempdir,
+  ]
+  common.RunAndCheckOutput(cmd, verbose=True)
+
+  logger.info("Writing vendor partitions to output archive.")
+  with zipfile.ZipFile(
+      output_zip_path, "a", compression=zipfile.ZIP_DEFLATED,
+      allowZip64=True) as output_zip:
+    for p in OPTIONS.vendor_partitions:
+      path = "IMAGES/{}.img".format(p)
+      common.ZipWrite(output_zip, os.path.join(vendor_tempdir, path), path)
+
+
 def main(argv):
 
   key_mapping_options = []
@@ -1289,6 +1352,10 @@
       OPTIONS.gki_signing_algorithm = a
     elif o == "--gki_signing_extra_args":
       OPTIONS.gki_signing_extra_args = a
+    elif o == "--vendor_otatools":
+      OPTIONS.vendor_otatools = a
+    elif o == "--vendor_partitions":
+      OPTIONS.vendor_partitions = set(a.split(","))
     else:
       return False
     return True
@@ -1339,6 +1406,8 @@
           "gki_signing_key=",
           "gki_signing_algorithm=",
           "gki_signing_extra_args=",
+          "vendor_partitions=",
+          "vendor_otatools=",
       ],
       extra_option_handler=option_handler)
 
@@ -1384,8 +1453,11 @@
   common.ZipClose(input_zip)
   common.ZipClose(output_zip)
 
+  if OPTIONS.vendor_partitions and OPTIONS.vendor_otatools:
+    BuildVendorPartitions(args[1])
+
   # Skip building userdata.img and cache.img when signing the target files.
-  new_args = ["--is_signing"]
+  new_args = ["--is_signing", "--add_missing", "--verbose"]
   # add_img_to_target_files builds the system image from scratch, so the
   # recovery patch is guaranteed to be regenerated there.
   if OPTIONS.rebuild_recovery: