releasetools: Support verifying AVB signed images with chained partitions.

For example, verify a target_files.zip that has system AVB-signed as a
chained partition.

  $ build/make/tools/releasetools/validate_target_files.py \
      signed-target_files-4904652.zip \
      --verity_key verifiedboot_pub.pem \
      --avb_system_key_path system_pub.pem

Note that verifiedboot_pub.pem should be the key (either public or
private) to verify vbmeta.img, and 'system_pub.pem' should be the key
(either public or private) for the chained partition of system.

testdata/testkey.key is the private key converted from
testdata/testkey.pk8 for testing purpose (`openssl pkcs8 -in
testdata/testkey.pk8 -inform DER -out testdata/testkey.key -nocrypt`).

Bug: 63706333
Test: python -m unittest test_common
Test: python -m unittest test_add_img_to_target_files
Test: `m dist` on aosp_walleye-userdebug; Run validate_target_files.py
      on the generated target_files.zip.
Test: Set up walleye with chained system partition; `m dist`; Run
      validate_target_files.py on the generated target_files.zip.
Change-Id: I38517ab39baf8a5bc1a6062fab2fe229b68e897d
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 5b5e6ed..c10e57c 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -371,6 +371,40 @@
     cmd.extend(["--salt", avb_salt])
 
 
+def GetAvbChainedPartitionArg(partition, info_dict, key=None):
+  """Constructs and returns the arg to build or verify a chained partition.
+
+  Args:
+    partition: The partition name.
+    info_dict: The info dict to look up the key info and rollback index
+        location.
+    key: The key to be used for building or verifying the partition. Defaults to
+        the key listed in info_dict.
+
+  Returns:
+    A string of form "partition:rollback_index_location:key" that can be used to
+    build or verify vbmeta image.
+
+  Raises:
+    AssertionError: When it fails to extract the public key with avbtool.
+  """
+  if key is None:
+    key = info_dict["avb_" + partition + "_key_path"]
+  avbtool = os.getenv('AVBTOOL') or info_dict["avb_avbtool"]
+  pubkey_path = MakeTempFile(prefix="avb-", suffix=".pubkey")
+  proc = Run(
+      [avbtool, "extract_public_key", "--key", key, "--output", pubkey_path],
+      stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+  stdoutdata, _ = proc.communicate()
+  assert proc.returncode == 0, \
+      "Failed to extract pubkey for {}:\n{}".format(
+          partition, stdoutdata)
+
+  rollback_index_location = info_dict[
+      "avb_" + partition + "_rollback_index_location"]
+  return "{}:{}:{}".format(partition, rollback_index_location, pubkey_path)
+
+
 def _BuildBootableImage(sourcedir, fs_config_file, info_dict=None,
                         has_ramdisk=False, two_step_image=False):
   """Build a bootable image from the specified sourcedir.