certify_bootimg: avoid using subprocess calls to unpack_bootimg

Uses function calls instead. This makes a certify_bootimg
binary can work without a unpack_bootimg binary. This can
reduce pre-built binaries in a kernel tree.

Bug: 180712476
Bug: 226121398
Test: atest --host mkbootimg_test
Test: atest --host certify_bootimg_test
Change-Id: I48c9f6767cb7b2c3e3bead008e630eeef310fca3
diff --git a/Android.bp b/Android.bp
index 77a5a32..65b6ac2 100644
--- a/Android.bp
+++ b/Android.bp
@@ -91,6 +91,20 @@
     ],
 }
 
+python_binary_host {
+    name: "certify_bootimg",
+    defaults: ["mkbootimg_defaults"],
+    main: "gki/certify_bootimg.py",
+    srcs: [
+        "gki/certify_bootimg.py",
+        "gki/generate_gki_certificate.py",
+        "unpack_bootimg.py",
+    ],
+    required: [
+        "avbtool",
+    ],
+}
+
 python_test_host {
     name: "mkbootimg_test",
     defaults: ["mkbootimg_defaults"],
diff --git a/gki/Android.bp b/gki/Android.bp
index 5173852..c62e7d8 100644
--- a/gki/Android.bp
+++ b/gki/Android.bp
@@ -16,20 +16,6 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
-python_binary_host {
-    name: "certify_bootimg",
-    defaults: ["mkbootimg_defaults"],
-    main: "certify_bootimg.py",
-    srcs: [
-        "certify_bootimg.py",
-        "generate_gki_certificate.py",
-    ],
-    required: [
-        "avbtool",
-        "unpack_bootimg",
-    ],
-}
-
 python_test_host {
     name: "certify_bootimg_test",
     defaults: ["mkbootimg_defaults"],
diff --git a/gki/certify_bootimg.py b/gki/certify_bootimg.py
index 5a99c1f..8067526 100755
--- a/gki/certify_bootimg.py
+++ b/gki/certify_bootimg.py
@@ -25,7 +25,8 @@
 import subprocess
 import tempfile
 
-from generate_gki_certificate import generate_gki_certificate
+from gki.generate_gki_certificate import generate_gki_certificate
+from unpack_bootimg import unpack_bootimg
 
 BOOT_SIGNATURE_SIZE = 16 * 1024
 
@@ -33,14 +34,7 @@
 def get_kernel(boot_img):
     """Extracts the kernel from |boot_img| and returns it."""
     with tempfile.TemporaryDirectory() as unpack_dir:
-        unpack_bootimg_cmd = [
-            'unpack_bootimg',
-            '--boot_img', boot_img,
-            '--out', unpack_dir,
-        ]
-        subprocess.run(unpack_bootimg_cmd, check=True,
-                       stdout=subprocess.DEVNULL)
-
+        unpack_bootimg(boot_img, unpack_dir)
         with open(os.path.join(unpack_dir, 'kernel'), 'rb') as kernel:
             kernel_bytes = kernel.read()
             assert len(kernel_bytes) > 0
diff --git a/unpack_bootimg.py b/unpack_bootimg.py
index 5c7d24c..462190f 100755
--- a/unpack_bootimg.py
+++ b/unpack_bootimg.py
@@ -19,7 +19,7 @@
 Extracts the kernel, ramdisk, second bootloader, dtb and recovery dtbo images.
 """
 
-from argparse import ArgumentParser, FileType, RawDescriptionHelpFormatter
+from argparse import ArgumentParser, RawDescriptionHelpFormatter
 from struct import unpack
 import os
 import shlex
@@ -181,12 +181,12 @@
         return args
 
 
-def unpack_boot_image(args):
+def unpack_boot_image(boot_img, output_dir):
     """extracts kernel, ramdisk, second bootloader and recovery dtbo"""
     info = BootImageInfoFormatter()
-    info.boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
+    info.boot_magic = unpack('8s', boot_img.read(8))[0].decode()
 
-    kernel_ramdisk_second_info = unpack('9I', args.boot_img.read(9 * 4))
+    kernel_ramdisk_second_info = unpack('9I', boot_img.read(9 * 4))
     # header_version is always at [8] regardless of the value of header_version.
     info.header_version = kernel_ramdisk_second_info[8]
 
@@ -199,7 +199,7 @@
         info.second_load_address = kernel_ramdisk_second_info[5]
         info.tags_load_address = kernel_ramdisk_second_info[6]
         info.page_size = kernel_ramdisk_second_info[7]
-        os_version_patch_level = unpack('I', args.boot_img.read(1 * 4))[0]
+        os_version_patch_level = unpack('I', boot_img.read(1 * 4))[0]
     else:
         info.kernel_size = kernel_ramdisk_second_info[0]
         info.ramdisk_size = kernel_ramdisk_second_info[1]
@@ -212,31 +212,31 @@
 
     if info.header_version < 3:
         info.product_name = cstr(unpack('16s',
-                                        args.boot_img.read(16))[0].decode())
-        info.cmdline = cstr(unpack('512s', args.boot_img.read(512))[0].decode())
-        args.boot_img.read(32)  # ignore SHA
+                                        boot_img.read(16))[0].decode())
+        info.cmdline = cstr(unpack('512s', boot_img.read(512))[0].decode())
+        boot_img.read(32)  # ignore SHA
         info.extra_cmdline = cstr(unpack('1024s',
-                                         args.boot_img.read(1024))[0].decode())
+                                         boot_img.read(1024))[0].decode())
     else:
         info.cmdline = cstr(unpack('1536s',
-                                   args.boot_img.read(1536))[0].decode())
+                                   boot_img.read(1536))[0].decode())
 
     if info.header_version in {1, 2}:
-        info.recovery_dtbo_size = unpack('I', args.boot_img.read(1 * 4))[0]
-        info.recovery_dtbo_offset = unpack('Q', args.boot_img.read(8))[0]
-        info.boot_header_size = unpack('I', args.boot_img.read(4))[0]
+        info.recovery_dtbo_size = unpack('I', boot_img.read(1 * 4))[0]
+        info.recovery_dtbo_offset = unpack('Q', boot_img.read(8))[0]
+        info.boot_header_size = unpack('I', boot_img.read(4))[0]
     else:
         info.recovery_dtbo_size = 0
 
     if info.header_version == 2:
-        info.dtb_size = unpack('I', args.boot_img.read(4))[0]
-        info.dtb_load_address = unpack('Q', args.boot_img.read(8))[0]
+        info.dtb_size = unpack('I', boot_img.read(4))[0]
+        info.dtb_load_address = unpack('Q', boot_img.read(8))[0]
     else:
         info.dtb_size = 0
         info.dtb_load_address = 0
 
     if info.header_version >= 4:
-        info.boot_signature_size = unpack('I', args.boot_img.read(4))[0]
+        info.boot_signature_size = unpack('I', boot_img.read(4))[0]
     else:
         info.boot_signature_size = 0
 
@@ -284,10 +284,10 @@
         image_info_list.append((boot_signature_offset, info.boot_signature_size,
                                 'boot_signature'))
 
-    create_out_dir(args.out)
+    create_out_dir(output_dir)
     for offset, size, name in image_info_list:
-        extract_image(offset, size, args.boot_img, os.path.join(args.out, name))
-    info.image_dir = args.out
+        extract_image(offset, size, boot_img, os.path.join(output_dir, name))
+    info.image_dir = output_dir
 
     return info
 
@@ -378,20 +378,20 @@
         return args
 
 
-def unpack_vendor_boot_image(args):
+def unpack_vendor_boot_image(boot_img, output_dir):
     info = VendorBootImageInfoFormatter()
-    info.boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
-    info.header_version = unpack('I', args.boot_img.read(4))[0]
-    info.page_size = unpack('I', args.boot_img.read(4))[0]
-    info.kernel_load_address = unpack('I', args.boot_img.read(4))[0]
-    info.ramdisk_load_address = unpack('I', args.boot_img.read(4))[0]
-    info.vendor_ramdisk_size = unpack('I', args.boot_img.read(4))[0]
-    info.cmdline = cstr(unpack('2048s', args.boot_img.read(2048))[0].decode())
-    info.tags_load_address = unpack('I', args.boot_img.read(4))[0]
-    info.product_name = cstr(unpack('16s', args.boot_img.read(16))[0].decode())
-    info.header_size = unpack('I', args.boot_img.read(4))[0]
-    info.dtb_size = unpack('I', args.boot_img.read(4))[0]
-    info.dtb_load_address = unpack('Q', args.boot_img.read(8))[0]
+    info.boot_magic = unpack('8s', boot_img.read(8))[0].decode()
+    info.header_version = unpack('I', boot_img.read(4))[0]
+    info.page_size = unpack('I', boot_img.read(4))[0]
+    info.kernel_load_address = unpack('I', boot_img.read(4))[0]
+    info.ramdisk_load_address = unpack('I', boot_img.read(4))[0]
+    info.vendor_ramdisk_size = unpack('I', boot_img.read(4))[0]
+    info.cmdline = cstr(unpack('2048s', boot_img.read(2048))[0].decode())
+    info.tags_load_address = unpack('I', boot_img.read(4))[0]
+    info.product_name = cstr(unpack('16s', boot_img.read(16))[0].decode())
+    info.header_size = unpack('I', boot_img.read(4))[0]
+    info.dtb_size = unpack('I', boot_img.read(4))[0]
+    info.dtb_load_address = unpack('Q', boot_img.read(8))[0]
 
     # Convenient shorthand.
     page_size = info.page_size
@@ -408,10 +408,10 @@
         (ramdisk_offset_base, info.vendor_ramdisk_size, 'vendor_ramdisk'))
 
     if info.header_version > 3:
-        info.vendor_ramdisk_table_size = unpack('I', args.boot_img.read(4))[0]
-        vendor_ramdisk_table_entry_num = unpack('I', args.boot_img.read(4))[0]
-        vendor_ramdisk_table_entry_size = unpack('I', args.boot_img.read(4))[0]
-        info.vendor_bootconfig_size = unpack('I', args.boot_img.read(4))[0]
+        info.vendor_ramdisk_table_size = unpack('I', boot_img.read(4))[0]
+        vendor_ramdisk_table_entry_num = unpack('I', boot_img.read(4))[0]
+        vendor_ramdisk_table_entry_size = unpack('I', boot_img.read(4))[0]
+        info.vendor_bootconfig_size = unpack('I', boot_img.read(4))[0]
         num_vendor_ramdisk_table_pages = get_number_of_pages(
             info.vendor_ramdisk_table_size, page_size)
         vendor_ramdisk_table_offset = page_size * (
@@ -422,16 +422,16 @@
         for idx in range(vendor_ramdisk_table_entry_num):
             entry_offset = vendor_ramdisk_table_offset + (
                 vendor_ramdisk_table_entry_size * idx)
-            args.boot_img.seek(entry_offset)
-            ramdisk_size = unpack('I', args.boot_img.read(4))[0]
-            ramdisk_offset = unpack('I', args.boot_img.read(4))[0]
-            ramdisk_type = unpack('I', args.boot_img.read(4))[0]
+            boot_img.seek(entry_offset)
+            ramdisk_size = unpack('I', boot_img.read(4))[0]
+            ramdisk_offset = unpack('I', boot_img.read(4))[0]
+            ramdisk_type = unpack('I', boot_img.read(4))[0]
             ramdisk_name = cstr(unpack(
                 f'{VENDOR_RAMDISK_NAME_SIZE}s',
-                args.boot_img.read(VENDOR_RAMDISK_NAME_SIZE))[0].decode())
+                boot_img.read(VENDOR_RAMDISK_NAME_SIZE))[0].decode())
             board_id = unpack(
                 f'{VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE}I',
-                args.boot_img.read(
+                boot_img.read(
                     4 * VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE))
             output_ramdisk_name = f'vendor_ramdisk{idx:02}'
 
@@ -455,14 +455,14 @@
     if info.dtb_size > 0:
         image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
 
-    create_out_dir(args.out)
+    create_out_dir(output_dir)
     for offset, size, name in image_info_list:
-        extract_image(offset, size, args.boot_img, os.path.join(args.out, name))
-    info.image_dir = args.out
+        extract_image(offset, size, boot_img, os.path.join(output_dir, name))
+    info.image_dir = output_dir
 
     if info.header_version > 3:
         vendor_ramdisk_by_name_dir = os.path.join(
-            args.out, 'vendor-ramdisk-by-name')
+            output_dir, 'vendor-ramdisk-by-name')
         create_out_dir(vendor_ramdisk_by_name_dir)
         for src, dst in vendor_ramdisk_symlinks:
             src_pathname = os.path.join('..', src)
@@ -475,19 +475,26 @@
     return info
 
 
-def unpack_image(args):
-    boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
-    args.boot_img.seek(0)
-    if boot_magic == 'ANDROID!':
-        info = unpack_boot_image(args)
-    elif boot_magic == 'VNDRBOOT':
-        info = unpack_vendor_boot_image(args)
-    else:
-        raise ValueError(f'Not an Android boot image, magic: {boot_magic}')
+def unpack_bootimg(boot_img, output_dir):
+    """Unpacks the |boot_img| to |output_dir|, and returns the 'info' object."""
+    with open(boot_img, 'rb') as image_file:
+        boot_magic = unpack('8s', image_file.read(8))[0].decode()
+        image_file.seek(0)
+        if boot_magic == 'ANDROID!':
+            info = unpack_boot_image(image_file, output_dir)
+        elif boot_magic == 'VNDRBOOT':
+            info = unpack_vendor_boot_image(image_file, output_dir)
+        else:
+            raise ValueError(f'Not an Android boot image, magic: {boot_magic}')
 
-    if args.format == 'mkbootimg':
+    return info
+
+
+def print_bootimg_info(info, output_format, null_separator):
+    """Format and print boot image info."""
+    if output_format == 'mkbootimg':
         mkbootimg_args = info.format_mkbootimg_argument()
-        if args.null:
+        if null_separator:
             print('\0'.join(mkbootimg_args) + '\0', end='')
         else:
             print(shlex.join(mkbootimg_args))
@@ -533,7 +540,7 @@
         description='Unpacks boot, recovery or vendor_boot image.',
         epilog=get_unpack_usage(),
     )
-    parser.add_argument('--boot_img', type=FileType('rb'), required=True,
+    parser.add_argument('--boot_img', required=True,
                         help='path to the boot, recovery or vendor_boot image')
     parser.add_argument('--out', default='out',
                         help='output directory of the unpacked images')
@@ -548,7 +555,8 @@
 def main():
     """parse arguments and unpack boot image"""
     args = parse_cmdline()
-    unpack_image(args)
+    info = unpack_bootimg(args.boot_img, args.out)
+    print_bootimg_info(info, args.format, args.null)
 
 
 if __name__ == '__main__':