Support re-generating DTBO image from add_img_to_target_files.py.

This is a step to enable signing a given target_files zip with release
keys.

When calling sign_target_files_apks.py, we will delete all the entries
under IMAGES/ in order to re-generate them (with the proper release
keys). In order to support that, we need to pack everything in need into
TF.zip.

Steps to test the CL.
a) Choose a target that has both AVB and DTBO enabled.
 $ m dist

b) Check IMAGES/dtbo.img and PREBUILT_IMAGES/dtbo.img both exist in the
   generated out/dist/TF.zip.

c) Remove the entries under IMAGES/ from the generated TF.zip.
 $ zip -d TF.zip IMAGES/\*

d) Re-generate the images with TF.zip.
 $ build/make/tools/releasetools/add_img_to_target_files.py TF.zip

e) Check that IMAGES/dtbo.img is re-generated, and it's identical to the
   image in b). Note that by default the re-generated image will carry a
   different footer, because of the random salt. This CL is verified by
   specifying the same salt.

Bug: 38315721
Test: see above.
Change-Id: I0bdc4e1cd4800962dc3902ca550dad6a8ca56c78
diff --git a/core/Makefile b/core/Makefile
index ac8b2f2..2a2be95 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -1781,7 +1781,7 @@
 
 ifdef INSTALLED_DTBOIMAGE_TARGET
 INTERNAL_AVB_MAKE_VBMETA_IMAGE_ARGS += \
-		--include_descriptors_from_image $(INSTALLED_DTBOIMAGE_TARGET)
+    --include_descriptors_from_image $(INSTALLED_DTBOIMAGE_TARGET)
 endif
 
 ifeq ($(BOARD_BUILD_SYSTEM_ROOT_IMAGE),true)
@@ -2292,8 +2292,14 @@
 	$(hide) cp $(INSTALLED_VENDORIMAGE_TARGET) $(zip_root)/IMAGES/
 endif
 ifdef BOARD_PREBUILT_DTBOIMAGE
-	$(hide) mkdir -p $(zip_root)/IMAGES
-	$(hide) cp $(INSTALLED_DTBOIMAGE_TARGET) $(zip_root)/IMAGES/
+	$(hide) mkdir -p $(zip_root)/PREBUILT_IMAGES
+	$(hide) cp $(INSTALLED_DTBOIMAGE_TARGET) $(zip_root)/PREBUILT_IMAGES/
+	$(hide) echo "has_dtbo=true" >> $(zip_root)/META/misc_info.txt
+ifeq ($(BOARD_AVB_ENABLE),true)
+	$(hide) echo "dtbo_size=$(BOARD_DTBOIMG_PARTITION_SIZE)" >> $(zip_root)/META/misc_info.txt
+	$(hide) echo "board_avb_dtbo_add_hash_footer_args=$(BOARD_AVB_DTBO_ADD_HASH_FOOTER_ARGS)" \
+	    >> $(zip_root)/META/misc_info.txt
+endif
 endif
 	@# Run fs_config on all the system, vendor, boot ramdisk,
 	@# and recovery ramdisk files in the zip, and save the output
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index 550a5fb..2f104cd 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -173,13 +173,43 @@
               block_list=block_list)
   return img.name
 
-def FindDtboPrebuilt(prefix="IMAGES/"):
-  """Find the prebuilt image of DTBO partition."""
 
-  prebuilt_path = os.path.join(OPTIONS.input_tmp, prefix, "dtbo.img")
-  if os.path.exists(prebuilt_path):
-    return prebuilt_path
-  return None
+def AddDtbo(output_zip, prefix="IMAGES/"):
+  """Adds the DTBO image.
+
+  Uses the image under prefix if it already exists. Otherwise looks for the
+  image under PREBUILT_IMAGES/, signs it as needed, and returns the image name.
+  """
+
+  img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "dtbo.img")
+  if os.path.exists(img.input_name):
+    print("dtbo.img already exists in %s, no need to rebuild..." % (prefix,))
+    return img.input_name
+
+  dtbo_prebuilt_path = os.path.join(
+      OPTIONS.input_tmp, "PREBUILT_IMAGES", "dtbo.img")
+  assert os.path.exists(dtbo_prebuilt_path)
+  shutil.copy(dtbo_prebuilt_path, img.name)
+
+  # AVB-sign the image as needed.
+  if OPTIONS.info_dict.get("board_avb_enable") == "true":
+    avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
+    part_size = OPTIONS.info_dict.get("dtbo_size")
+    # The AVB hash footer will be replaced if already present.
+    cmd = [avbtool, "add_hash_footer", "--image", img.name,
+           "--partition_size", str(part_size), "--partition_name", "dtbo"]
+    common.AppendAVBSigningArgs(cmd)
+    args = OPTIONS.info_dict.get("board_avb_dtbo_add_hash_footer_args")
+    if args and args.strip():
+      cmd.extend(shlex.split(args))
+    p = common.Run(cmd, stdout=subprocess.PIPE)
+    p.communicate()
+    assert p.returncode == 0, \
+        "avbtool add_hash_footer of %s failed" % (img.name,)
+
+  img.Write()
+  return img.name
+
 
 def CreateImage(input_dir, info_dict, what, output_file, block_list=None):
   print("creating " + what + ".img...")
@@ -303,7 +333,7 @@
               dtbo_img_path, prefix="IMAGES/"):
   """Create a VBMeta image and store it in output_zip."""
   img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "vbmeta.img")
-  avbtool = os.getenv('AVBTOOL') or "avbtool"
+  avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
   cmd = [avbtool, "make_vbmeta_image",
          "--output", img.name,
          "--include_descriptors_from_image", boot_img_path,
@@ -312,10 +342,10 @@
     cmd.extend(["--include_descriptors_from_image", vendor_img_path])
   if dtbo_img_path is not None:
     cmd.extend(["--include_descriptors_from_image", dtbo_img_path])
-  if OPTIONS.info_dict.get("system_root_image", None) == "true":
+  if OPTIONS.info_dict.get("system_root_image") == "true":
     cmd.extend(["--setup_rootfs_from_kernel", system_img_path])
   common.AppendAVBSigningArgs(cmd)
-  args = OPTIONS.info_dict.get("board_avb_make_vbmeta_image_args", None)
+  args = OPTIONS.info_dict.get("board_avb_make_vbmeta_image_args")
   if args and args.strip():
     cmd.extend(shlex.split(args))
   p = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
@@ -495,7 +525,7 @@
 
   banner("system")
   system_img_path = AddSystem(
-    output_zip, recovery_img=recovery_image, boot_img=boot_image)
+      output_zip, recovery_img=recovery_image, boot_img=boot_image)
   vendor_img_path = None
   if has_vendor:
     banner("vendor")
@@ -508,13 +538,19 @@
     AddUserdata(output_zip)
     banner("cache")
     AddCache(output_zip)
-  if OPTIONS.info_dict.get("board_bpt_enable", None) == "true":
+
+  if OPTIONS.info_dict.get("board_bpt_enable") == "true":
     banner("partition-table")
     AddPartitionTable(output_zip)
-  if OPTIONS.info_dict.get("board_avb_enable", None) == "true":
+
+  dtbo_img_path = None
+  if OPTIONS.info_dict.get("has_dtbo") == "true":
+    banner("dtbo")
+    dtbo_img_path = AddDtbo(output_zip)
+
+  if OPTIONS.info_dict.get("board_avb_enable") == "true":
     banner("vbmeta")
     boot_contents = boot_image.WriteToTemp()
-    dtbo_img_path = FindDtboPrebuilt()
     AddVBMeta(output_zip, boot_contents.name, system_img_path,
               vendor_img_path, dtbo_img_path)