Merge "Write the vbmeta digest to the output zipfile also" into sc-dev

tools/releasetools/add_img_to_target_files.py

@@ -704,6 +704,31 @@ def AddApexInfo(output_zip):
       common.ZipWrite(output_zip, output_file, arc_name)
 
 
+def AddVbmetaDigest(output_zip):
+  """Write the vbmeta digest to the output dir and zipfile."""
+
+  # Calculate the vbmeta digest and put the result in to META/
+  boot_images = OPTIONS.info_dict.get("boot_images")
+  # Disable the digest calculation if the target_file is used as a container
+  # for boot images.
+  boot_container = boot_images and len(boot_images.split()) >= 2
+  if (OPTIONS.info_dict.get("avb_enable") == "true" and not boot_container and
+      OPTIONS.info_dict.get("avb_building_vbmeta_image") == "true"):
+    avbtool = OPTIONS.info_dict["avb_avbtool"]
+    digest = verity_utils.CalculateVbmetaDigest(OPTIONS.input_tmp, avbtool)
+    vbmeta_digest_txt = os.path.join(OPTIONS.input_tmp, "META",
+                                     "vbmeta_digest.txt")
+    with open(vbmeta_digest_txt, 'w') as f:
+      f.write(digest)
+    # writes to the output zipfile
+    if output_zip:
+      arc_name = "META/vbmeta_digest.txt"
+      if arc_name in output_zip.namelist():
+        OPTIONS.replace_updated_files_list.append(arc_name)
+      else:
+        common.ZipWriteStr(output_zip, arc_name, digest)
+
+
 def AddImagesToTargetFiles(filename):
   """Creates and adds images (boot/recovery/system/...) to a target_files.zip.
 
@@ -957,19 +982,7 @@ def AddImagesToTargetFiles(filename):
     with open(pack_radioimages_txt) as f:
       AddPackRadioImages(output_zip, f.readlines())
 
-  # Calculate the vbmeta digest and put the result in to META/
+  AddVbmetaDigest(output_zip)
-  boot_images = OPTIONS.info_dict.get("boot_images")
-  # Disable the digest calculation if the target_file is used as a container
-  # for boot images.
-  boot_container = boot_images and len(boot_images.split()) >= 2
-  if (OPTIONS.info_dict.get("avb_enable") == "true" and not boot_container and
-      OPTIONS.info_dict.get("avb_building_vbmeta_image") == "true"):
-    avbtool = OPTIONS.info_dict["avb_avbtool"]
-    digest = verity_utils.CalculateVbmetaDigest(OPTIONS.input_tmp, avbtool)
-    vbmeta_digest_txt = os.path.join(OPTIONS.input_tmp, "META",
-                                     "vbmeta_digest.txt")
-    with open(vbmeta_digest_txt, 'w') as f:
-      f.write(digest)
 
   if output_zip:
     common.ZipClose(output_zip)