Add pvmfw partition to target files

The pvmfw is included in an `m` build but not in the distribution files.
Apply the same treatment as the dtbo partition to copy the pvmfw
partition to dist/ since, similar to the dtbo image, the pvmfw image is
generally provided as a prebuilt image.

Test: make dist
Bug: 174457787
Change-Id: I6f42517ba42db92e90048d1236d7255ccbd73f73
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index 5f9f19a..900c7b5 100644
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -350,6 +350,41 @@
   img.Write()
   return img.name
 
+def AddPvmfw(output_zip):
+  """Adds the pvmfw image.
+
+  Uses the image under IMAGES/ 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, "IMAGES", "pvmfw.img")
+  if os.path.exists(img.name):
+    logger.info("pvmfw.img already exists; no need to rebuild...")
+    return img.name
+
+  pvmfw_prebuilt_path = os.path.join(
+      OPTIONS.input_tmp, "PREBUILT_IMAGES", "pvmfw.img")
+  assert os.path.exists(pvmfw_prebuilt_path)
+  shutil.copy(pvmfw_prebuilt_path, img.name)
+
+  # AVB-sign the image as needed.
+  if OPTIONS.info_dict.get("avb_enable") == "true":
+    # Signing requires +w
+    os.chmod(img.name, os.stat(img.name).st_mode | stat.S_IWUSR)
+
+    avbtool = OPTIONS.info_dict["avb_avbtool"]
+    part_size = OPTIONS.info_dict["pvmfw_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", "pvmfw"]
+    common.AppendAVBSigningArgs(cmd, "pvmfw")
+    args = OPTIONS.info_dict.get("avb_pvmfw_add_hash_footer_args")
+    if args and args.strip():
+      cmd.extend(shlex.split(args))
+    common.RunAndCheckOutput(cmd)
+
+  img.Write()
+  return img.name
+
 def AddCustomImages(output_zip, partition_name):
   """Adds and signs custom images in IMAGES/.
 
@@ -948,6 +983,10 @@
     banner("dtbo")
     partitions['dtbo'] = AddDtbo(output_zip)
 
+  if OPTIONS.info_dict.get("has_pvmfw") == "true":
+    banner("pvmfw")
+    partitions['pvmfw'] = AddPvmfw(output_zip)
+
   # Custom images.
   custom_partitions = OPTIONS.info_dict.get(
       "avb_custom_images_partition_list", "").strip().split()