certify_bootimg: support more archive formats

The common archive format in a kernel tree is usually *.tar.gz
rather than *.zip. Supports different archive formats for
--boot_img_archive (renamed from --boot_img_zip).

An usage example:
  certify_bootimg --boot_img_archive boot-img.tar.gz \
      --algorithm SHA256_RSA4096 \
      --key external/avb/test/data/testkey_rsa4096.pem \
      --extra_args "--prop foo:bar" \
      --extra_args "--prop gki:nice" \
      --output boot-certified-img.tar.gz

The formats of the input archive and the output archive can
be different:
  certify_bootimg --boot_img_archive boot-img.zip \
      --algorithm SHA256_RSA4096 \
      --key external/avb/test/data/testkey_rsa4096.pem \
      --extra_args "--prop foo:bar" \
      --extra_args "--prop gki:nice" \
      --output boot-certified-img.tar

The supported archive formats in the python binary built by
the current Android build system are:
  shutil.get_unpack_formats():
    [('gztar', ['.tar.gz', '.tgz'], "gzip'ed tar-file"),
     ('tar', ['.tar'], 'uncompressed tar file'),
     ('zip', ['.zip'], 'ZIP file')]

Bug: 223288963
Test: atest --host certify_bootimg_test
Change-Id: I74542c435a7b16f66708530e1c00d6fa8331e4cd
diff --git a/gki/certify_bootimg.py b/gki/certify_bootimg.py
index 9a7b058..8e3d1af 100755
--- a/gki/certify_bootimg.py
+++ b/gki/certify_bootimg.py
@@ -160,6 +160,20 @@
     return d
 
 
+def get_archive_name_and_format_for_shutil(path):
+    """Returns archive name and format to shutil.make_archive() for the |path|.
+
+    e.g., returns ('/path/to/boot-img', 'gztar') if |path| is
+    '/path/to/boot-img.tar.gz'.
+    """
+    for format_name, format_extensions, _ in shutil.get_unpack_formats():
+        for extension in format_extensions:
+            if path.endswith(extension):
+                return path[:-len(extension)], format_name
+
+    raise ValueError(f"Unsupported archive format: '{path}'")
+
+
 def parse_cmdline():
     """Parse command-line options."""
     parser = ArgumentParser(add_help=True)
@@ -169,7 +183,7 @@
     input_group.add_argument(
         '--boot_img', help='path to the boot image to certify')
     input_group.add_argument(
-        '--boot_img_zip', help='path to the boot-img-*.zip archive to certify')
+        '--boot_img_archive', help='path to the boot images archive to certify')
 
     parser.add_argument('--algorithm', required=True,
                         help='signing algorithm for the certificate')
@@ -208,35 +222,45 @@
         shutil.copy2(boot_tmp, output_img)
 
 
-def certify_bootimg_zip(boot_img_zip, output_zip, algorithm, key, extra_args):
-    """Similar to certify_bootimg(), but for a zip archive of boot images."""
-    with tempfile.TemporaryDirectory() as unzip_dir:
-        shutil.unpack_archive(boot_img_zip, unzip_dir)
+def certify_bootimg_archive(boot_img_archive, output_archive,
+                            algorithm, key, extra_args):
+    """Similar to certify_bootimg(), but for an archive of boot images."""
+    with tempfile.TemporaryDirectory() as unpack_dir:
+        shutil.unpack_archive(boot_img_archive, unpack_dir)
 
-        gki_info_file = os.path.join(unzip_dir, 'gki-info.txt')
+        gki_info_file = os.path.join(unpack_dir, 'gki-info.txt')
         if os.path.exists(gki_info_file):
             info_dict = load_dict_from_file(gki_info_file)
             if 'certify_bootimg_extra_args' in info_dict:
                 extra_args.extend(
                     shlex.split(info_dict['certify_bootimg_extra_args']))
 
-        for boot_img in glob.glob(os.path.join(unzip_dir, 'boot-*.img')):
+        for boot_img in glob.glob(os.path.join(unpack_dir, 'boot-*.img')):
             print(f'Certifying {os.path.basename(boot_img)} ...')
             certify_bootimg(boot_img=boot_img, output_img=boot_img,
                             algorithm=algorithm, key=key, extra_args=extra_args)
 
-        print(f'Making certified archive: {output_zip}')
-        archive_base_name = os.path.splitext(output_zip)[0]
-        shutil.make_archive(archive_base_name, 'zip', unzip_dir)
+        print(f'Making certified archive: {output_archive}')
+        archive_file_name, archive_format = (
+            get_archive_name_and_format_for_shutil(output_archive))
+        built_archive = shutil.make_archive(archive_file_name,
+                                            archive_format,
+                                            unpack_dir)
+        # shutil.make_archive() builds *.tar.gz when then |archive_format| is
+        # 'gztar'. However, the end user might specify |output_archive| with
+        # *.tgz. Renaming *.tar.gz to *.tgz for this case.
+        if built_archive != os.path.realpath(output_archive):
+            print(f'Renaming {built_archive} -> {output_archive} ...')
+            os.rename(built_archive, output_archive)
 
 
 def main():
     """Parse arguments and certify the boot image."""
     args = parse_cmdline()
 
-    if args.boot_img_zip:
-        certify_bootimg_zip(args.boot_img_zip, args.output, args.algorithm,
-                            args.key, args.extra_args)
+    if args.boot_img_archive:
+        certify_bootimg_archive(args.boot_img_archive, args.output,
+                                args.algorithm, args.key, args.extra_args)
     else:
         certify_bootimg(args.boot_img, args.output, args.algorithm,
                         args.key, args.extra_args)
diff --git a/gki/certify_bootimg_test.py b/gki/certify_bootimg_test.py
index 8c7c4d3..b7d367e 100644
--- a/gki/certify_bootimg_test.py
+++ b/gki/certify_bootimg_test.py
@@ -68,20 +68,27 @@
         subprocess.check_call(avbtool_cmd)
 
 
-def generate_test_boot_image_archive(output_zip, boot_img_info, gki_info=None):
-    """Generates a zip archive of test boot images.
+def generate_test_boot_image_archive(archive_file_name, archive_format,
+                                     boot_img_info, gki_info=None):
+    """Generates an archive of test boot images.
 
     It also adds a file gki-info.txt, which contains additional settings for
     for `certify_bootimg --extra_args`.
 
     Args:
-        output_zip: the output zip archive, e.g., /path/to/boot-img.zip.
+        archive_file_name: the name of the archive file to create, including the
+          path, minus any format-specific extension.
+        archive_format: the |format| parameter for shutil.make_archive().
+          e.g., 'zip', 'tar', or 'gztar', etc.
         boot_img_info: a list of (boot_image_name, kernel_size,
           partition_size) tuples. e.g.,
           [('boot-1.0.img', 4096, 4 * 1024),
            ('boot-2.0.img', 8192, 8 * 1024)].
         gki_info: the file content to be written into 'gki-info.txt' in the
-          |output_zip|.
+          created archive.
+
+    Returns:
+        The full path of the created archive. e.g., /path/to/boot-img.tar.gz.
     """
     with tempfile.TemporaryDirectory() as temp_out_dir:
         for name, kernel_size, partition_size in boot_img_info:
@@ -96,8 +103,9 @@
             with open(gki_info_path, 'w', encoding='utf-8') as f:
                 f.write(gki_info)
 
-        archive_base_name = os.path.splitext(output_zip)[0]
-        shutil.make_archive(archive_base_name, 'zip', temp_out_dir)
+        return shutil.make_archive(archive_file_name,
+                                   archive_format,
+                                   temp_out_dir)
 
 
 def has_avb_footer(image):
@@ -175,10 +183,10 @@
         boot_signature_bytes = boot_signature_bytes[next_signature_size:]
 
 
-def extract_boot_archive_with_signatures(boot_img_zip, output_dir):
+def extract_boot_archive_with_signatures(boot_img_archive, output_dir):
     """Extracts boot images and signatures of a boot images archive.
 
-    Suppose there are two boot images in |boot_img_zip|: boot-1.0.img
+    Suppose there are two boot images in |boot_img_archive|: boot-1.0.img
     and boot-2.0.img. This function then extracts each boot-*.img and
     their signatures as:
       - |output_dir|/boot-1.0.img
@@ -188,7 +196,7 @@
       - |output_dir|/boot-2.0/boot_signature1
       - |output_dir|/boot-2.0/boot_signature2
     """
-    shutil.unpack_archive(boot_img_zip, output_dir)
+    shutil.unpack_archive(boot_img_archive, output_dir)
     for boot_img in glob.glob(os.path.join(output_dir, 'boot-*.img')):
         img_name = os.path.splitext(os.path.basename(boot_img))[0]
         signature_output_dir = os.path.join(output_dir, img_name)
@@ -647,37 +655,38 @@
                               err.stderr)
 
     def test_certify_bootimg_archive(self):
-        """Tests certify_bootimg for a boot-img.zip."""
+        """Tests certify_bootimg for a boot images archive.."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
-            boot_img_zip = os.path.join(temp_out_dir, 'boot-img.zip')
+            boot_img_archive_name = os.path.join(temp_out_dir, 'boot-img')
             gki_info = ('certify_bootimg_extra_args='
                         '--prop KERNEL_RELEASE:5.10.42'
                         '-android13-0-00544-ged21d463f856 '
                         '--prop BRANCH:android13-5.10-2022-05 '
                         '--prop BUILD_NUMBER:ab8295296 '
                         '--prop SPACE:"nice to meet you"\n')
-            generate_test_boot_image_archive(
-                boot_img_zip,
+            boot_img_archive_path = generate_test_boot_image_archive(
+                boot_img_archive_name,
+                'gztar',
                 # A list of (boot_img_name, kernel_size, partition_size).
                 [('boot-1.0.img', 8 * 1024, 128 * 1024),
                  ('boot-2.0.img', 16 * 1024, 256 * 1024)],
                 gki_info)
 
             # Certify the boot image archive, with a RSA4096 key.
-            boot_certified_img_zip = os.path.join(temp_out_dir,
-                                                  'boot-certified-img.zip')
+            boot_certified_img_archive = os.path.join(
+                temp_out_dir, 'boot-certified-img.tar.gz')
             certify_bootimg_cmds = [
                 'certify_bootimg',
-                '--boot_img_zip', boot_img_zip,
+                '--boot_img_archive', boot_img_archive_path,
                 '--algorithm', 'SHA256_RSA4096',
                 '--key', './testdata/testkey_rsa4096.pem',
                 '--extra_args', '--prop gki:nice '
                 '--prop space:"nice to meet you"',
-                '--output', boot_certified_img_zip,
+                '--output', boot_certified_img_archive,
             ]
             subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
 
-            extract_boot_archive_with_signatures(boot_certified_img_zip,
+            extract_boot_archive_with_signatures(boot_certified_img_archive,
                                                  temp_out_dir)
 
             # Checks an AVB footer exists and the image size remains.
@@ -701,54 +710,56 @@
                     self._EXPECTED_BOOT_2_0_SIGNATURE2_RSA4096})
 
     def test_certify_bootimg_archive_without_gki_info(self):
-        """Tests certify_bootimg for a boot-img.zip."""
+        """Tests certify_bootimg for a boot images archive."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
-            boot_img_zip = os.path.join(temp_out_dir, 'boot-img.zip')
+            boot_img_archive_name = os.path.join(temp_out_dir, 'boot-img')
 
-            # Checks ceritfy_bootimg works for a boot-img.zip without a
-            # gki-info.txt.
-            generate_test_boot_image_archive(
-                boot_img_zip,
+            # Checks ceritfy_bootimg works for a boot images archive without a
+            # gki-info.txt. Using *.zip -> *.tar.
+            boot_img_archive_path = generate_test_boot_image_archive(
+                boot_img_archive_name,
+                'zip',
                 # A list of (boot_img_name, kernel_size, partition_size).
                 [('boot-3.0.img', 8 * 1024, 128 * 1024)],
                 gki_info=None)
             # Certify the boot image archive, with a RSA4096 key.
-            boot_certified_img_zip = os.path.join(temp_out_dir,
-                                                  'boot-certified-img.zip')
+            boot_certified_img_archive = os.path.join(
+                temp_out_dir, 'boot-certified-img.tar')
             certify_bootimg_cmds = [
                 'certify_bootimg',
-                '--boot_img_zip', boot_img_zip,
+                '--boot_img_archive', boot_img_archive_path,
                 '--algorithm', 'SHA256_RSA4096',
                 '--key', './testdata/testkey_rsa4096.pem',
                 '--extra_args', '--prop gki:nice '
                 '--prop space:"nice to meet you"',
-                '--output', boot_certified_img_zip,
+                '--output', boot_certified_img_archive,
             ]
             subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
 
-            # Checks ceritfy_bootimg works for a boot-img.zip with a special
-            # gki-info.txt.
-            generate_test_boot_image_archive(
-                boot_img_zip,
+            # Checks ceritfy_bootimg works for a boot images archive with a
+            # special gki-info.txt. Using *.tar -> *.tgz.
+            boot_img_archive_path = generate_test_boot_image_archive(
+                boot_img_archive_name,
+                'tar',
                 # A list of (boot_img_name, kernel_size, partition_size).
                 [('boot-3.0.img', 8 * 1024, 128 * 1024)],
                 gki_info='a=b\n'
                          'c=d\n')
             # Certify the boot image archive, with a RSA4096 key.
-            boot_certified_img_zip = os.path.join(temp_out_dir,
-                                                  'boot-certified-img.zip')
+            boot_certified_img_archive2 = os.path.join(
+                temp_out_dir, 'boot-certified-img.tgz')
             certify_bootimg_cmds = [
                 'certify_bootimg',
-                '--boot_img_zip', boot_img_zip,
+                '--boot_img_archive', boot_img_archive_path,
                 '--algorithm', 'SHA256_RSA4096',
                 '--key', './testdata/testkey_rsa4096.pem',
                 '--extra_args', '--prop gki:nice '
                 '--prop space:"nice to meet you"',
-                '--output', boot_certified_img_zip,
+                '--output', boot_certified_img_archive2,
             ]
             subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
 
-            extract_boot_archive_with_signatures(boot_certified_img_zip,
+            extract_boot_archive_with_signatures(boot_certified_img_archive2,
                                                  temp_out_dir)
 
             # Checks an AVB footer exists and the image size remains.