Merge "Attempt to run mkbootimg_test as part of unit tests CI" am: 45e0eea0dc am: d19d0b9dae am: 6329cde2af am: 0691fd1fef am: 768504755e

Original change: https://android-review.googlesource.com/c/platform/system/tools/mkbootimg/+/1664783

Change-Id: Ie053307f1467554c75c2dd7f9f351d6c7b4af3f8
diff --git a/repack_bootimg.py b/repack_bootimg.py
index 4cc63db..c320018 100755
--- a/repack_bootimg.py
+++ b/repack_bootimg.py
@@ -21,9 +21,11 @@
 """
 
 import argparse
+import datetime
 import enum
-import json
+import glob
 import os
+import shlex
 import shutil
 import subprocess
 import tempfile
@@ -75,22 +77,28 @@
     """Enum class for different boot image types."""
     BOOT_IMAGE = 1
     VENDOR_BOOT_IMAGE = 2
+    SINGLE_RAMDISK_FRAGMENT = 3
+    MULTIPLE_RAMDISK_FRAGMENTS = 4
 
 
 class RamdiskImage:
     """A class that supports packing/unpacking a ramdisk."""
-    def __init__(self, ramdisk_img):
+    def __init__(self, ramdisk_img, unpack=True):
         self._ramdisk_img = ramdisk_img
         self._ramdisk_format = None
         self._ramdisk_dir = None
         self._temp_file_manager = TempFileManager()
 
-        self._unpack_ramdisk()
+        if unpack:
+            self._unpack_ramdisk()
+        else:
+            self._ramdisk_dir = self._temp_file_manager.make_temp_dir(
+                suffix='_new_ramdisk')
 
     def _unpack_ramdisk(self):
         """Unpacks the ramdisk."""
         self._ramdisk_dir = self._temp_file_manager.make_temp_dir(
-            suffix=os.path.basename(self._ramdisk_img))
+            suffix='_' + os.path.basename(self._ramdisk_img))
 
         # The compression format might be in 'lz4' or 'gzip' format,
         # trying lz4 first.
@@ -115,18 +123,10 @@
             # toybox cpio arguments:
             #   -i: extract files from stdin
             #   -d: create directories if needed
-            cpio_result = subprocess.run(
-                ['toybox', 'cpio', '-id'], check=False,
-                input=decompressed_result.stdout, capture_output=True,
-                cwd=self._ramdisk_dir)
-
-            # toybox cpio command might return a non-zero code, e.g., found
-            # duplicated files in the ramdisk. Treat it as non-fatal with
-            # check=False and only print the error message here.
-            if cpio_result.returncode != 0:
-                print('\n'
-                      'WARNNING: cpio command error:\n' +
-                      cpio_result.stderr.decode('utf-8').strip() + '\n')
+            #   -u: override existing files
+            subprocess.run(
+                ['toybox', 'cpio', '-idu'], check=True,
+                input=decompressed_result.stdout, cwd=self._ramdisk_dir)
 
             print("=== Unpacked ramdisk: '{}' ===".format(
                 self._ramdisk_img))
@@ -168,64 +168,55 @@
         self._bootimg_dir = None
         self._bootimg_type = None
         self._ramdisk = None
-        # Potential images to extract from a boot.img. Unlike the ramdisk,
-        # the content of the following images will not be changed during the
-        # repack process.
-        self._intact_image_candidates = ('dtb', 'kernel',
-                                         'recovery_dtbo', 'second')
-        self._repack_intact_image_args = []
+        self._previous_mkbootimg_args = []
         self._temp_file_manager = TempFileManager()
 
         self._unpack_bootimg()
 
+    def _get_vendor_ramdisks(self):
+        """Returns a list of vendor ramdisks after unpack."""
+        return sorted(glob.glob(
+            os.path.join(self._bootimg_dir, 'vendor_ramdisk*')))
+
     def _unpack_bootimg(self):
         """Unpacks the boot.img and the ramdisk inside."""
         self._bootimg_dir = self._temp_file_manager.make_temp_dir(
-            suffix=os.path.basename(self._bootimg))
+            suffix='_' + os.path.basename(self._bootimg))
 
         # Unpacks the boot.img first.
-        subprocess.check_call(
-            ['unpack_bootimg', '--boot_img', self._bootimg,
-             '--out', self._bootimg_dir], stdout=subprocess.DEVNULL)
+        unpack_bootimg_cmds = [
+            'unpack_bootimg',
+            '--boot_img', self._bootimg,
+            '--out', self._bootimg_dir,
+            '--format=mkbootimg',
+        ]
+        result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                capture_output=True, encoding='utf-8')
+        self._previous_mkbootimg_args = shlex.split(result.stdout)
         print("=== Unpacked boot image: '{}' ===".format(self._bootimg))
 
-        for img_name in self._intact_image_candidates:
-            img_file = os.path.join(self._bootimg_dir, img_name)
-            if os.path.exists(img_file):
-                # Prepares args for repacking those intact images. e.g.,
-                # --kernel kernel_image, --dtb dtb_image.
-                self._repack_intact_image_args.extend(
-                    ['--' + img_name, img_file])
-
         # From the output dir, checks there is 'ramdisk' or 'vendor_ramdisk'.
         ramdisk = os.path.join(self._bootimg_dir, 'ramdisk')
         vendor_ramdisk = os.path.join(self._bootimg_dir, 'vendor_ramdisk')
+        vendor_ramdisks = self._get_vendor_ramdisks()
         if os.path.exists(ramdisk):
             self._ramdisk = RamdiskImage(ramdisk)
             self._bootimg_type = BootImageType.BOOT_IMAGE
         elif os.path.exists(vendor_ramdisk):
             self._ramdisk = RamdiskImage(vendor_ramdisk)
             self._bootimg_type = BootImageType.VENDOR_BOOT_IMAGE
+        elif len(vendor_ramdisks) == 1:
+            self._ramdisk = RamdiskImage(vendor_ramdisks[0])
+            self._bootimg_type = BootImageType.SINGLE_RAMDISK_FRAGMENT
+        elif len(vendor_ramdisks) > 1:
+            # Creates an empty RamdiskImage() below, without unpack.
+            # We'll then add files into this newly created ramdisk, then pack
+            # it with other vendor ramdisks together.
+            self._ramdisk = RamdiskImage(ramdisk_img=None, unpack=False)
+            self._bootimg_type = BootImageType.MULTIPLE_RAMDISK_FRAGMENTS
         else:
             raise RuntimeError('Both ramdisk and vendor_ramdisk do not exist.')
 
-    @property
-    def _previous_mkbootimg_args(self):
-        """Returns the previous mkbootimg args from mkbootimg_args.json file."""
-        # Loads the saved mkbootimg_args.json from previous unpack_bootimg.
-        command = []
-        mkbootimg_config = os.path.join(
-            self._bootimg_dir, 'mkbootimg_args.json')
-        with open (mkbootimg_config) as config:
-            mkbootimg_args = json.load(config)
-            for argname, value in mkbootimg_args.items():
-                # argname, e.g., 'board', 'header_version', etc., does not have
-                # prefix '--', which is required when invoking `mkbootimg.py`.
-                # Prepends '--' to make the full args, e.g., --header_version.
-                command.extend(['--' + argname, value])
-        return command
-
-
     def repack_bootimg(self):
         """Repacks the ramdisk and rebuild the boot.img"""
 
@@ -238,39 +229,59 @@
         # Uses previous mkbootimg args, e.g., --vendor_cmdline, --dtb_offset.
         mkbootimg_cmd.extend(self._previous_mkbootimg_args)
 
-        if self._repack_intact_image_args:
-            mkbootimg_cmd.extend(self._repack_intact_image_args)
-
-        if self._bootimg_type == BootImageType.VENDOR_BOOT_IMAGE:
-            mkbootimg_cmd.extend(['--vendor_ramdisk', new_ramdisk])
-            mkbootimg_cmd.extend(['--vendor_boot', self._bootimg])
-            # TODO(bowgotsai): add support for multiple vendor ramdisk.
-        else:
-            mkbootimg_cmd.extend(['--ramdisk', new_ramdisk])
+        ramdisk_option = ''
+        if self._bootimg_type == BootImageType.BOOT_IMAGE:
+            ramdisk_option = '--ramdisk'
             mkbootimg_cmd.extend(['--output', self._bootimg])
+        elif self._bootimg_type == BootImageType.VENDOR_BOOT_IMAGE:
+            ramdisk_option = '--vendor_ramdisk'
+            mkbootimg_cmd.extend(['--vendor_boot', self._bootimg])
+        elif self._bootimg_type == BootImageType.SINGLE_RAMDISK_FRAGMENT:
+            ramdisk_option = '--vendor_ramdisk_fragment'
+            mkbootimg_cmd.extend(['--vendor_boot', self._bootimg])
+        elif self._bootimg_type == BootImageType.MULTIPLE_RAMDISK_FRAGMENTS:
+            mkbootimg_cmd.extend(['--ramdisk_type', 'PLATFORM'])
+            ramdisk_name = (
+                'RAMDISK_' +
+                datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S'))
+            mkbootimg_cmd.extend(['--ramdisk_name', ramdisk_name])
+            mkbootimg_cmd.extend(['--vendor_ramdisk_fragment', new_ramdisk])
+            mkbootimg_cmd.extend(['--vendor_boot', self._bootimg])
+
+        if ramdisk_option and ramdisk_option not in mkbootimg_cmd:
+            raise RuntimeError("Failed to find '{}' from:\n  {}".format(
+                ramdisk_option, shlex.join(mkbootimg_cmd)))
+        # Replaces the original ramdisk with the newly packed ramdisk.
+        if ramdisk_option:
+            ramdisk_index = mkbootimg_cmd.index(ramdisk_option) + 1
+            mkbootimg_cmd[ramdisk_index] = new_ramdisk
 
         subprocess.check_call(mkbootimg_cmd)
         print("=== Repacked boot image: '{}' ===".format(self._bootimg))
 
-
     def add_files(self, src_dir, files):
         """Copy files from the src_dir into current ramdisk.
 
         Args:
             src_dir: a source dir containing the files to copy from.
-            files: a list of files to copy from src_dir.
+            files: a list of files or src_file:dst_file pairs to copy from
+              src_dir to the current ramdisk.
         """
         # Creates missing parent dirs with 0o755.
         original_mask = os.umask(0o022)
         for f in files:
-            src_file = os.path.join(src_dir, f)
-            dst_file = os.path.join(self.ramdisk_dir, f)
+            if ':' in f:
+                src_file = os.path.join(src_dir, f.split(':')[0])
+                dst_file = os.path.join(self.ramdisk_dir, f.split(':')[1])
+            else:
+                src_file = os.path.join(src_dir, f)
+                dst_file = os.path.join(self.ramdisk_dir, f)
+
             dst_dir = os.path.dirname(dst_file)
             if not os.path.exists(dst_dir):
                 print("Creating dir '{}'".format(dst_dir))
                 os.makedirs(dst_dir, 0o755)
-            print("Copying file '{}' into '{}'".format(
-                src_file, self._bootimg))
+            print("Copying file '{}' into '{}'".format(src_file, dst_file))
             shutil.copy2(src_file, dst_file)
         os.umask(original_mask)
 
@@ -280,9 +291,47 @@
         return self._ramdisk.ramdisk_dir
 
 
+def _get_repack_usage():
+    return """Usage examples:
+
+  * --ramdisk_add
+
+    Specifies a list of files or src_file:dst_file pairs to copy from
+    --src_bootimg's ramdisk into --dst_bootimg's ramdisk.
+
+    $ repack_bootimg \\
+        --src_bootimg boot-debug-5.4.img --dst_bootimg vendor_boot-debug.img \\
+        --ramdisk_add first_stage_ramdisk/userdebug_plat_sepolicy.cil:userdebug_plat_sepolicy.cil
+
+    The above command copies '/first_stage_ramdisk/userdebug_plat_sepolicy.cil'
+    from --src_bootimg's ramdisk to '/userdebug_plat_sepolicy.cil' of
+    --dst_bootimg's ramdisk, then repacks the --dst_bootimg.
+
+    $ repack_bootimg \\
+        --src_bootimg boot-debug-5.4.img --dst_bootimg vendor_boot-debug.img \\
+        --ramdisk_add first_stage_ramdisk/userdebug_plat_sepolicy.cil
+
+    This is similar to the previous example, but the source file path and
+    destination file path are the same:
+        '/first_stage_ramdisk/userdebug_plat_sepolicy.cil'.
+
+    We can also combine both usage together with a list of copy instructions.
+    For example:
+
+    $ repack_bootimg \\
+        --src_bootimg boot-debug-5.4.img --dst_bootimg vendor_boot-debug.img \\
+        --ramdisk_add file1 file2:/subdir/file2 file3
+"""
+
+
 def _parse_args():
     """Parse command-line options."""
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description='Repacks boot, recovery or vendor_boot image by importing'
+                    'ramdisk files from --src_bootimg to --dst_bootimg.',
+        epilog=_get_repack_usage(),
+    )
 
     parser.add_argument(
         '--src_bootimg', help='filename to source boot image',
@@ -291,8 +340,10 @@
         '--dst_bootimg', help='filename to destination boot image',
         type=str, required=True)
     parser.add_argument(
-        '--ramdisk_add', help='a list of files to add into the ramdisk',
-        nargs='+', default=['first_stage_ramdisk/userdebug_plat_sepolicy.cil']
+        '--ramdisk_add', nargs='+',
+        help='a list of files or src_file:dst_file pairs to add into '
+             'the ramdisk',
+        default=['userdebug_plat_sepolicy.cil']
     )
 
     return parser.parse_args()
diff --git a/tests/mkbootimg_test.py b/tests/mkbootimg_test.py
index 110b3cb..ae5cf6b 100644
--- a/tests/mkbootimg_test.py
+++ b/tests/mkbootimg_test.py
@@ -17,7 +17,6 @@
 """Tests mkbootimg and unpack_bootimg."""
 
 import filecmp
-import json
 import logging
 import os
 import random
@@ -35,9 +34,13 @@
 VENDOR_BOOT_ARGS_OFFSET = 28
 VENDOR_BOOT_ARGS_SIZE = 2048
 
-
 BOOT_IMAGE_V4_SIGNATURE_SIZE = 4096
 
+TEST_KERNEL_CMDLINE = (
+    'printk.devkmsg=on firmware_class.path=/vendor/etc/ init=/init '
+    'kfence.sample_interval=500 loop.max_part=7 bootconfig'
+)
+
 
 def generate_test_file(pathname, size, seed=None):
     """Generates a gibberish-filled test file and returns its pathname."""
@@ -96,7 +99,7 @@
                 '--header_version', '4',
                 '--kernel', kernel,
                 '--ramdisk', ramdisk,
-                '--cmdline', 'test-cmdline',
+                '--cmdline', TEST_KERNEL_CMDLINE,
                 '--os_version', '11.0.0',
                 '--os_patch_level', '2021-01',
                 '--gki_signing_algorithm', 'SHA256_RSA2048',
@@ -141,8 +144,8 @@
                 '      Partition Name:        boot\n'
                 '      Salt:                  d00df00d\n'
                 '      Digest:                '
-                '9749bb508f2677426b14ff668d39a163'
-                'e16f0c4cbaf92ec096124e3f199fafac\n'
+                'cf3755630856f23ab70e501900050fee'
+                'f30b633b3e82a9085a578617e344f9c7\n'
                 '      Flags:                 0\n'
                 "    Prop: foo -> 'bar'\n"
                 "    Prop: gki -> 'nice'\n"
@@ -180,7 +183,7 @@
                 '--header_version', '4',
                 '--kernel', kernel,
                 '--ramdisk', ramdisk,
-                '--cmdline', 'test-cmdline',
+                '--cmdline', TEST_KERNEL_CMDLINE,
                 '--os_version', '11.0.0',
                 '--os_patch_level', '2021-01',
                 '--gki_signing_avbtool_path', self._avbtool_path,
@@ -218,7 +221,7 @@
                 '--header_version', '4',
                 '--kernel', kernel,
                 '--ramdisk', ramdisk,
-                '--cmdline', 'test-cmdline',
+                '--cmdline', TEST_KERNEL_CMDLINE,
                 '--os_version', '11.0.0',
                 '--os_patch_level', '2021-01',
                 '--output', boot_img,
@@ -263,6 +266,7 @@
                 '--board_id0', '0xC0FFEE',
                 '--board_id15', '0x15151515',
                 '--vendor_ramdisk_fragment', ramdisk2,
+                '--vendor_cmdline', TEST_KERNEL_CMDLINE,
                 '--vendor_bootconfig', bootconfig,
             ]
             unpack_bootimg_cmds = [
@@ -274,6 +278,7 @@
                 'boot magic: VNDRBOOT',
                 'vendor boot image header version: 4',
                 'vendor ramdisk total size: 16384',
+                f'vendor command line args: {TEST_KERNEL_CMDLINE}',
                 'dtb size: 4096',
                 'vendor ramdisk table size: 324',
                 'size: 4096', 'offset: 0', 'type: 0x1', 'name:',
@@ -309,12 +314,12 @@
                 ])
                 self.fail(msg)
 
-    def test_unpack_vendor_boot_image_v4_format_mkbootimg(self):
-        """Tests `unpack_bootimg --format=mkbootimg`."""
+    def test_unpack_vendor_boot_image_v4(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
-            vendor_boot_reconstructed = os.path.join(
-                temp_out_dir, 'vendor_boot.reconstructed')
+            vendor_boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'vendor_boot.img.reconstructed')
             dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
             ramdisk1 = generate_test_file(
                 os.path.join(temp_out_dir, 'ramdisk1'), 0x121212)
@@ -337,6 +342,7 @@
                 '--board_id0', '0xC0FFEE',
                 '--board_id15', '0x15151515',
                 '--vendor_ramdisk_fragment', ramdisk2,
+                '--vendor_cmdline', TEST_KERNEL_CMDLINE,
                 '--vendor_bootconfig', bootconfig,
             ]
             unpack_bootimg_cmds = [
@@ -350,14 +356,14 @@
                                     capture_output=True, encoding='utf-8')
             mkbootimg_cmds = [
                 'mkbootimg',
-                '--vendor_boot', vendor_boot_reconstructed,
+                '--vendor_boot', vendor_boot_img_reconstructed,
             ]
             unpack_format_args = shlex.split(result.stdout)
             mkbootimg_cmds.extend(unpack_format_args)
 
             subprocess.run(mkbootimg_cmds, check=True)
             self.assertTrue(
-                filecmp.cmp(vendor_boot_img, vendor_boot_reconstructed),
+                filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
                 'reconstructed vendor_boot image differ from the original')
 
             # Also check that -0, --null are as expected.
@@ -368,100 +374,22 @@
             self.assertEqual('\0'.join(unpack_format_args) + '\0',
                              unpack_format_null_args)
 
-    def test_unpack_vendor_boot_image_v3_format_mkbootimg(self):
-        """Tests `unpack_bootimg --format=mkbootimg`."""
+    def test_unpack_vendor_boot_image_v3(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
-            vendor_boot_reconstructed = os.path.join(
-                temp_out_dir, 'vendor_boot.reconstructed')
-            dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
-            ramdisk1 = generate_test_file(
-                os.path.join(temp_out_dir, 'ramdisk1'), 0x121212)
-
-            mkbootimg_cmds = [
-                'mkbootimg',
-                '--header_version', '3',
-                '--vendor_boot', vendor_boot_img,
-                '--dtb', dtb,
-                '--vendor_ramdisk', ramdisk1,
-            ]
-            unpack_bootimg_cmds = [
-                'unpack_bootimg',
-                '--boot_img', vendor_boot_img,
-                '--out', os.path.join(temp_out_dir, 'out'),
-                '--format=mkbootimg',
-            ]
-            subprocess.run(mkbootimg_cmds, check=True)
-            result = subprocess.run(unpack_bootimg_cmds, check=True,
-                                    capture_output=True, encoding='utf-8')
-            mkbootimg_cmds = [
-                'mkbootimg',
-                '--vendor_boot', vendor_boot_reconstructed,
-            ]
-            mkbootimg_cmds.extend(shlex.split(result.stdout))
-
-            subprocess.run(mkbootimg_cmds, check=True)
-            self.assertTrue(
-                filecmp.cmp(vendor_boot_img, vendor_boot_reconstructed),
-                'reconstructed vendor_boot image differ from the original')
-
-    def test_unpack_boot_image_v3_json_args(self):
-        """Tests mkbootimg_args.json when unpacking a boot image version 3."""
-        with tempfile.TemporaryDirectory() as temp_out_dir:
-            boot_img = os.path.join(temp_out_dir, 'boot.img')
-            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
-                                        0x1000)
-            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
-                                         0x1000)
-            mkbootimg_cmds = [
-                'mkbootimg',
-                '--header_version', '3',
-                '--kernel', kernel,
-                '--ramdisk', ramdisk,
-                '--cmdline', 'test-cmdline',
-                '--os_version', '11.0.0',
-                '--os_patch_level', '2021-01',
-                '--output', boot_img,
-            ]
-            unpack_bootimg_cmds = [
-                'unpack_bootimg',
-                '--boot_img', boot_img,
-                '--out', os.path.join(temp_out_dir, 'out'),
-            ]
-            # The expected dict in mkbootimg_args.json.
-            expected_mkbootimg_args = {
-                'cmdline': 'test-cmdline',
-                'header_version': '3',
-                'os_patch_level': '2021-01',
-                'os_version': '11.0.0'
-            }
-
-            subprocess.run(mkbootimg_cmds, check=True)
-            subprocess.run(unpack_bootimg_cmds, check=True)
-
-            json_file = os.path.join(
-                temp_out_dir, 'out', 'mkbootimg_args.json')
-            with open(json_file) as json_fd:
-                actual_mkbootimg_args = json.load(json_fd)
-                self.assertEqual(actual_mkbootimg_args,
-                                 expected_mkbootimg_args)
-
-    def test_unpack_vendor_boot_image_v3_json_args(self):
-        """Tests mkbootimg_args.json when unpacking a vendor boot image version
-        3.
-        """
-        with tempfile.TemporaryDirectory() as temp_out_dir:
-            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
+            vendor_boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'vendor_boot.img.reconstructed')
             dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
             ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
-                                         0x1000)
+                                         0x121212)
             mkbootimg_cmds = [
                 'mkbootimg',
                 '--header_version', '3',
                 '--vendor_boot', vendor_boot_img,
                 '--vendor_ramdisk', ramdisk,
                 '--dtb', dtb,
-                '--vendor_cmdline', 'test-vendor_cmdline',
+                '--vendor_cmdline', TEST_KERNEL_CMDLINE,
                 '--board', 'product_name',
                 '--base', '0x00000000',
                 '--dtb_offset', '0x01f00000',
@@ -474,35 +402,70 @@
                 'unpack_bootimg',
                 '--boot_img', vendor_boot_img,
                 '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
             ]
-            # The expected dict in mkbootimg_args.json.
-            expected_mkbootimg_args = {
-                'header_version': '3',
-                'vendor_cmdline': 'test-vendor_cmdline',
-                'board': 'product_name',
-                'base': '0x00000000',
-                'dtb_offset': '0x0000000001f00000',  # dtb_offset is uint64_t.
-                'kernel_offset': '0x00008000',
-                'pagesize': '0x00001000',
-                'ramdisk_offset': '0x01000000',
-                'tags_offset': '0x00000100',
-            }
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--vendor_boot', vendor_boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
 
             subprocess.run(mkbootimg_cmds, check=True)
-            subprocess.run(unpack_bootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
+                'reconstructed vendor_boot image differ from the original')
 
-            json_file = os.path.join(
-                temp_out_dir, 'out', 'mkbootimg_args.json')
-            with open(json_file) as json_fd:
-                actual_mkbootimg_args = json.load(json_fd)
-                self.assertEqual(actual_mkbootimg_args,
-                                 expected_mkbootimg_args)
+    def test_unpack_boot_image_v3(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img = os.path.join(temp_out_dir, 'boot.img')
+            boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'boot.img.reconstructed')
+            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
+                                        0x1000)
+            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
+                                         0x1000)
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '3',
+                '--kernel', kernel,
+                '--ramdisk', ramdisk,
+                '--cmdline', TEST_KERNEL_CMDLINE,
+                '--os_version', '11.0.0',
+                '--os_patch_level', '2021-01',
+                '--output', boot_img,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
+            ]
 
-    def test_unpack_boot_image_v2_json_args(self):
-        """Tests mkbootimg_args.json when unpacking a boot image v2."""
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--out', boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(boot_img, boot_img_reconstructed),
+                'reconstructed boot image differ from the original')
+
+    def test_unpack_boot_image_v2(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             # Output image path.
             boot_img = os.path.join(temp_out_dir, 'boot.img')
+            boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'boot.img.reconstructed')
             # Creates blank images first.
             kernel = generate_test_file(
                 os.path.join(temp_out_dir, 'kernel'), 0x1000)
@@ -543,38 +506,30 @@
                 'unpack_bootimg',
                 '--boot_img', boot_img,
                 '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
             ]
-            # The expected dict in mkbootimg_args.json.
-            expected_mkbootimg_args = {
-                'header_version': '2',
-                'base': '0x00000000',
-                'kernel_offset': '0x00008000',
-                'ramdisk_offset': '0x01000000',
-                'second_offset': '0x40000000',
-                'dtb_offset': '0x0000000001f00000',  # dtb_offset is uint64_t.
-                'tags_offset': '0x00000100',
-                'pagesize': '0x00001000',
-                'os_version': '11.0.0',
-                'os_patch_level': '2021-03',
-                'board': 'boot_v2',
-                'cmdline': cmdline + extra_cmdline,
-            }
 
             subprocess.run(mkbootimg_cmds, check=True)
-            subprocess.run(unpack_bootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--out', boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
 
-            json_file = os.path.join(
-                temp_out_dir, 'out', 'mkbootimg_args.json')
-            with open(json_file) as json_fd:
-                actual_mkbootimg_args = json.load(json_fd)
-                self.assertEqual(actual_mkbootimg_args,
-                                 expected_mkbootimg_args)
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(boot_img, boot_img_reconstructed),
+                'reconstructed boot image differ from the original')
 
-    def test_unpack_boot_image_v1_json_args(self):
-        """Tests mkbootimg_args.json when unpacking a boot image v1."""
+    def test_unpack_boot_image_v1(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             # Output image path.
             boot_img = os.path.join(temp_out_dir, 'boot.img')
+            boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'boot.img.reconstructed')
             # Creates blank images first.
             kernel = generate_test_file(
                 os.path.join(temp_out_dir, 'kernel'), 0x1000)
@@ -607,37 +562,30 @@
                 'unpack_bootimg',
                 '--boot_img', boot_img,
                 '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
             ]
-            # The expected dict in mkbootimg_args.json.
-            expected_mkbootimg_args = {
-                'header_version': '1',
-                'base': '0x00000000',
-                'kernel_offset': '0x00008000',
-                'ramdisk_offset': '0x01000000',
-                'second_offset': '0x00000000',
-                'tags_offset': '0x00000100',
-                'pagesize': '0x00001000',
-                'os_version': '11.0.0',
-                'os_patch_level': '2021-03',
-                'board': 'boot_v1',
-                'cmdline': cmdline + extra_cmdline,
-            }
 
             subprocess.run(mkbootimg_cmds, check=True)
-            subprocess.run(unpack_bootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--out', boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
 
-            json_file = os.path.join(
-                temp_out_dir, 'out', 'mkbootimg_args.json')
-            with open(json_file) as json_fd:
-                actual_mkbootimg_args = json.load(json_fd)
-                self.assertEqual(actual_mkbootimg_args,
-                                 expected_mkbootimg_args)
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(boot_img, boot_img_reconstructed),
+                'reconstructed boot image differ from the original')
 
-    def test_unpack_boot_image_v0_json_args(self):
-        """Tests mkbootimg_args.json when unpacking a boot image v0."""
+    def test_unpack_boot_image_v0(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             # Output image path.
             boot_img = os.path.join(temp_out_dir, 'boot.img')
+            boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'boot.img.reconstructed')
             # Creates blank images first.
             kernel = generate_test_file(
                 os.path.join(temp_out_dir, 'kernel'), 0x1000)
@@ -672,30 +620,26 @@
                 '--boot_img', boot_img,
                 '--out', os.path.join(temp_out_dir, 'out'),
             ]
-            # The expected dict in mkbootimg_args.json.
-            expected_mkbootimg_args = {
-                'header_version': '0',
-                'base': '0x00000000',
-                'kernel_offset': '0x00008000',
-                'ramdisk_offset': '0x01000000',
-                'second_offset': '0x40000000',
-                'tags_offset': '0x00000100',
-                'pagesize': '0x00001000',
-                'os_version': '11.0.0',
-                'os_patch_level': '2021-03',
-                'board': 'boot_v0',
-                'cmdline': cmdline + extra_cmdline,
-            }
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
+            ]
 
             subprocess.run(mkbootimg_cmds, check=True)
-            subprocess.run(unpack_bootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--out', boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
 
-            json_file = os.path.join(
-                temp_out_dir, 'out', 'mkbootimg_args.json')
-            with open(json_file) as json_fd:
-                actual_mkbootimg_args = json.load(json_fd)
-                self.assertEqual(actual_mkbootimg_args,
-                                 expected_mkbootimg_args)
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(boot_img, boot_img_reconstructed),
+                'reconstructed boot image differ from the original')
 
     def test_boot_image_v2_cmdline_null_terminator(self):
         """Tests that kernel commandline is null-terminated."""
diff --git a/unpack_bootimg.py b/unpack_bootimg.py
index 7209137..2b176e5 100755
--- a/unpack_bootimg.py
+++ b/unpack_bootimg.py
@@ -21,7 +21,6 @@
 
 from argparse import ArgumentParser, FileType, RawDescriptionHelpFormatter
 from struct import unpack
-import json
 import os
 import shlex
 
@@ -29,8 +28,6 @@
 VENDOR_RAMDISK_NAME_SIZE = 32
 VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE = 16
 
-MKBOOTIMG_ARGS_FILE = 'mkbootimg_args.json'
-
 
 def create_out_dir(dir_path):
     """creates a directory 'dir_path' if it does not exist"""
@@ -69,13 +66,6 @@
     return '{:04d}-{:02d}'.format(y, m)
 
 
-def print_os_version_patch_level(value):
-    os_version = value >> 11
-    os_patch_level = value & ((1<<11) - 1)
-    print('os version: %s' % format_os_version(os_version))
-    print('os patch level: %s' % format_os_patch_level(os_patch_level))
-
-
 def decode_os_version_patch_level(os_version_patch_level):
     """Returns a tuple of (os_version, os_patch_level)."""
     os_version = os_version_patch_level >> 11
@@ -84,197 +74,216 @@
             format_os_patch_level(os_patch_level))
 
 
-def get_boot_image_v2_and_below_args(header_version, page_size,
-                                     kernel_load_address, ramdisk_load_address,
-                                     second_load_address, tags_load_address,
-                                     dtb_load_address, cmdline, extra_cmdline,
-                                     os_version_patch_level, product_name):
-    """Returns a dict of mkbootimg.py arguments for v0, v1 and v2 boot.img."""
-    mkbootimg_args = {}
-    mkbootimg_args['header_version'] = str(header_version)
-    # The type of pagesize is uint32_t, using '0xFFFFFFFF' as the output format.
-    mkbootimg_args['pagesize'] = '{:#010x}'.format(page_size)
+class BootImageInfoFormatter:
+    """Formats the boot image info."""
 
-    # Kernel load address is base + kernel_offset in mkbootimg.py.
-    # However, we don't know the value of 'base' when unpack a boot.img
-    # in this script. So always set 'base' to be zero and 'kernel_offset' to
-    # be the kernel load address. Same for 'ramdisk_offset', 'second_offset',
-    # 'tags_offset' and 'dtb_offset'.
-    # The following types are uint32_t, using '0xFFFFFFFF' as the output format.
-    mkbootimg_args['base'] = '{:#010x}'.format(0)
-    mkbootimg_args['kernel_offset'] = '{:#010x}'.format(kernel_load_address)
-    mkbootimg_args['ramdisk_offset'] = '{:#010x}'.format(ramdisk_load_address)
-    mkbootimg_args['second_offset'] = '{:#010x}'.format(second_load_address)
-    mkbootimg_args['tags_offset'] = '{:#010x}'.format(tags_load_address)
+    def format_pretty_text(self):
+        lines = []
+        lines.append(f'boot magic: {self.boot_magic}')
 
-    # dtb is added in boot image v2, and is absent in v1 or v0.
-    if header_version == 2:
-        # The type of dtb_offset is uint64_t, using '0xFFFFFFFFEEEEEEEE' as
-        # the output format.
-        mkbootimg_args['dtb_offset'] = '{:#018x}'.format(dtb_load_address)
+        if self.header_version < 3:
+            lines.append(f'kernel_size: {self.kernel_size}')
+            lines.append(
+                f'kernel load address: {self.kernel_load_address:#010x}')
+            lines.append(f'ramdisk size: {self.ramdisk_size}')
+            lines.append(
+                f'ramdisk load address: {self.ramdisk_load_address:#010x}')
+            lines.append(f'second bootloader size: {self.second_size}')
+            lines.append(
+                f'second bootloader load address: '
+                f'{self.second_load_address:#010x}')
+            lines.append(
+                f'kernel tags load address: {self.tags_load_address:#010x}')
+            lines.append(f'page size: {self.page_size}')
+        else:
+            lines.append(f'kernel_size: {self.kernel_size}')
+            lines.append(f'ramdisk size: {self.ramdisk_size}')
 
-    mkbootimg_args['os_version'], mkbootimg_args['os_patch_level'] = (
-        decode_os_version_patch_level(os_version_patch_level))
+        lines.append(f'os version: {self.os_version}')
+        lines.append(f'os patch level: {self.os_patch_level}')
+        lines.append(f'boot image header version: {self.header_version}')
 
-    mkbootimg_args['cmdline'] = cmdline + extra_cmdline
-    mkbootimg_args['board'] = product_name
+        if self.header_version < 3:
+            lines.append(f'product name: {self.product_name}')
 
-    return mkbootimg_args
+        lines.append(f'command line args: {self.cmdline}')
 
+        if self.header_version < 3:
+            lines.append(f'additional command line args: {self.extra_cmdline}')
 
-def get_boot_image_v3_args(header_version, os_version_patch_level, cmdline):
-    """Returns a dict of arguments to be used in mkbootimg.py later."""
-    mkbootimg_args = {}
-    mkbootimg_args['header_version'] = str(header_version)
-    mkbootimg_args['os_version'], mkbootimg_args['os_patch_level'] = (
-        decode_os_version_patch_level(os_version_patch_level))
-    mkbootimg_args['cmdline'] = cmdline
+        if self.header_version in {1, 2}:
+            lines.append(f'recovery dtbo size: {self.recovery_dtbo_size}')
+            lines.append(
+                f'recovery dtbo offset: {self.recovery_dtbo_offset:#018x}')
+            lines.append(f'boot header size: {self.boot_header_size}')
 
-    return mkbootimg_args
+        if self.header_version == 2:
+            lines.append(f'dtb size: {self.dtb_size}')
+            lines.append(f'dtb address: {self.dtb_load_address:#018x}')
+
+        if self.header_version >= 4:
+            lines.append(
+                f'boot.img signature size: {self.boot_signature_size}')
+
+        return '\n'.join(lines)
+
+    def format_mkbootimg_argument(self):
+        args = []
+        args.extend(['--header_version', str(self.header_version)])
+        args.extend(['--os_version', self.os_version])
+        args.extend(['--os_patch_level', self.os_patch_level])
+
+        args.extend(['--kernel', os.path.join(self.image_dir, 'kernel')])
+        args.extend(['--ramdisk', os.path.join(self.image_dir, 'ramdisk')])
+
+        if self.header_version <= 2:
+            if self.second_size > 0:
+                args.extend(['--second',
+                             os.path.join(self.image_dir, 'second')])
+            if self.recovery_dtbo_size > 0:
+                args.extend(['--recovery_dtbo',
+                             os.path.join(self.image_dir, 'recovery_dtbo')])
+            if self.dtb_size > 0:
+                args.extend(['--dtb', os.path.join(self.image_dir, 'dtb')])
+
+            args.extend(['--pagesize', f'{self.page_size:#010x}'])
+
+            # Kernel load address is base + kernel_offset in mkbootimg.py.
+            # However we don't know the value of 'base' when unpacking a boot
+            # image in this script, so we set 'base' to zero and 'kernel_offset'
+            # to the kernel load address, 'ramdisk_offset' to the ramdisk load
+            # address, ... etc.
+            args.extend(['--base', f'{0:#010x}'])
+            args.extend(['--kernel_offset',
+                         f'{self.kernel_load_address:#010x}'])
+            args.extend(['--ramdisk_offset',
+                         f'{self.ramdisk_load_address:#010x}'])
+            args.extend(['--second_offset',
+                         f'{self.second_load_address:#010x}'])
+            args.extend(['--tags_offset', f'{self.tags_load_address:#010x}'])
+
+            # dtb is added in boot image v2, and is absent in v1 or v0.
+            if self.header_version == 2:
+                # dtb_offset is uint64_t.
+                args.extend(['--dtb_offset', f'{self.dtb_load_address:#018x}'])
+
+            args.extend(['--board', self.product_name])
+            args.extend(['--cmdline', self.cmdline + self.extra_cmdline])
+        else:
+            args.extend(['--cmdline', self.cmdline])
+
+        return args
 
 
 def unpack_boot_image(args):
     """extracts kernel, ramdisk, second bootloader and recovery dtbo"""
-    boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
-    print('boot_magic: %s' % boot_magic)
-    # TODO(yochiang): Support --format=mkbootimg
+    info = BootImageInfoFormatter()
+    info.boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
 
     kernel_ramdisk_second_info = unpack('9I', args.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]
 
-    # version is always at [8] regardless of version.
-    version = kernel_ramdisk_second_info[8]
-
-    if version < 3:
-        kernel_size = kernel_ramdisk_second_info[0]
-        kernel_load_address = kernel_ramdisk_second_info[1]
-        ramdisk_size = kernel_ramdisk_second_info[2]
-        ramdisk_load_address = kernel_ramdisk_second_info[3]
-        second_size = kernel_ramdisk_second_info[4]
-        second_load_address = kernel_ramdisk_second_info[5]
-        tags_load_address = kernel_ramdisk_second_info[6]
-        page_size = kernel_ramdisk_second_info[7]
+    if info.header_version < 3:
+        info.kernel_size = kernel_ramdisk_second_info[0]
+        info.kernel_load_address = kernel_ramdisk_second_info[1]
+        info.ramdisk_size = kernel_ramdisk_second_info[2]
+        info.ramdisk_load_address = kernel_ramdisk_second_info[3]
+        info.second_size = kernel_ramdisk_second_info[4]
+        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]
     else:
-        kernel_size = kernel_ramdisk_second_info[0]
-        ramdisk_size = kernel_ramdisk_second_info[1]
+        info.kernel_size = kernel_ramdisk_second_info[0]
+        info.ramdisk_size = kernel_ramdisk_second_info[1]
         os_version_patch_level = kernel_ramdisk_second_info[2]
-        second_size = 0
-        page_size = BOOT_IMAGE_HEADER_V3_PAGESIZE
+        info.second_size = 0
+        info.page_size = BOOT_IMAGE_HEADER_V3_PAGESIZE
 
-    if version < 3:
-        print('kernel_size: %s' % kernel_size)
-        print('kernel load address: %#x' % kernel_load_address)
-        print('ramdisk size: %s' % ramdisk_size)
-        print('ramdisk load address: %#x' % ramdisk_load_address)
-        print('second bootloader size: %s' % second_size)
-        print('second bootloader load address: %#x' % second_load_address)
-        print('kernel tags load address: %#x' % tags_load_address)
-        print('page size: %s' % page_size)
-        print_os_version_patch_level(os_version_patch_level)
-    else:
-        print('kernel_size: %s' % kernel_size)
-        print('ramdisk size: %s' % ramdisk_size)
-        print_os_version_patch_level(os_version_patch_level)
+    info.os_version, info.os_patch_level = decode_os_version_patch_level(
+        os_version_patch_level)
 
-    print('boot image header version: %s' % version)
-
-    if version < 3:
-        product_name = cstr(unpack('16s', args.boot_img.read(16))[0].decode())
-        print('product name: %s' % product_name)
-        cmdline = cstr(unpack('512s', args.boot_img.read(512))[0].decode())
-        print('command line args: %s' % cmdline)
-    else:
-        cmdline = cstr(unpack('1536s', args.boot_img.read(1536))[0].decode())
-        print('command line args: %s' % cmdline)
-
-    if version < 3:
+    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
-
-    if version < 3:
-        extra_cmdline = cstr(unpack('1024s',
-                                    args.boot_img.read(1024))[0].decode())
-        print('additional command line args: %s' % extra_cmdline)
-
-    if 0 < version < 3:
-        recovery_dtbo_size = unpack('I', args.boot_img.read(1 * 4))[0]
-        print('recovery dtbo size: %s' % recovery_dtbo_size)
-        recovery_dtbo_offset = unpack('Q', args.boot_img.read(8))[0]
-        print('recovery dtbo offset: %#x' % recovery_dtbo_offset)
-        boot_header_size = unpack('I', args.boot_img.read(4))[0]
-        print('boot header size: %s' % boot_header_size)
+        info.extra_cmdline = cstr(unpack('1024s',
+                                         args.boot_img.read(1024))[0].decode())
     else:
-        recovery_dtbo_size = 0
+        info.cmdline = cstr(unpack('1536s',
+                                   args.boot_img.read(1536))[0].decode())
 
-    if 1 < version < 3:
-        dtb_size = unpack('I', args.boot_img.read(4))[0]
-        print('dtb size: %s' % dtb_size)
-        dtb_load_address = unpack('Q', args.boot_img.read(8))[0]
-        print('dtb address: %#x' % dtb_load_address)
+    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]
     else:
-        dtb_size = 0
-        dtb_load_address = 0
+        info.recovery_dtbo_size = 0
 
-    # Saves the arguments to be reused in mkbootimg.py later.
-    if version < 3:
-        mkbootimg_args = get_boot_image_v2_and_below_args(
-            version, page_size, kernel_load_address, ramdisk_load_address,
-            second_load_address, tags_load_address, dtb_load_address, cmdline,
-            extra_cmdline, os_version_patch_level, product_name)
+    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]
     else:
-        mkbootimg_args = get_boot_image_v3_args(
-            version, os_version_patch_level, cmdline)
-    with open(os.path.join(args.out, MKBOOTIMG_ARGS_FILE), 'w') as f:
-        json.dump(mkbootimg_args, f, sort_keys=True, indent=4)
+        info.dtb_size = 0
+        info.dtb_load_address = 0
 
-    if version >= 4:
-        boot_signature_size = unpack('I', args.boot_img.read(4))[0]
-        print('boot.img signature size: %s' % boot_signature_size)
+    if info.header_version >= 4:
+        info.boot_signature_size = unpack('I', args.boot_img.read(4))[0]
     else:
-        boot_signature_size = 0
+        info.boot_signature_size = 0
 
     # The first page contains the boot header
     num_header_pages = 1
 
-    num_kernel_pages = get_number_of_pages(kernel_size, page_size)
-    kernel_offset = page_size * num_header_pages  # header occupies a page
-    image_info_list = [(kernel_offset, kernel_size, 'kernel')]
+    # Convenient shorthand.
+    page_size = info.page_size
 
-    num_ramdisk_pages = get_number_of_pages(ramdisk_size, page_size)
+    num_kernel_pages = get_number_of_pages(info.kernel_size, page_size)
+    kernel_offset = page_size * num_header_pages  # header occupies a page
+    image_info_list = [(kernel_offset, info.kernel_size, 'kernel')]
+
+    num_ramdisk_pages = get_number_of_pages(info.ramdisk_size, page_size)
     ramdisk_offset = page_size * (num_header_pages + num_kernel_pages
                                  ) # header + kernel
-    image_info_list.append((ramdisk_offset, ramdisk_size, 'ramdisk'))
+    image_info_list.append((ramdisk_offset, info.ramdisk_size, 'ramdisk'))
 
-    if second_size > 0:
+    if info.second_size > 0:
         second_offset = page_size * (
             num_header_pages + num_kernel_pages + num_ramdisk_pages
             )  # header + kernel + ramdisk
-        image_info_list.append((second_offset, second_size, 'second'))
+        image_info_list.append((second_offset, info.second_size, 'second'))
 
-    if recovery_dtbo_size > 0:
-        image_info_list.append((recovery_dtbo_offset, recovery_dtbo_size,
+    if info.recovery_dtbo_size > 0:
+        image_info_list.append((info.recovery_dtbo_offset,
+                                info.recovery_dtbo_size,
                                 'recovery_dtbo'))
-    if dtb_size > 0:
-        num_second_pages = get_number_of_pages(second_size, page_size)
+    if info.dtb_size > 0:
+        num_second_pages = get_number_of_pages(info.second_size, page_size)
         num_recovery_dtbo_pages = get_number_of_pages(
-            recovery_dtbo_size, page_size)
+            info.recovery_dtbo_size, page_size)
         dtb_offset = page_size * (
             num_header_pages + num_kernel_pages + num_ramdisk_pages +
             num_second_pages + num_recovery_dtbo_pages)
 
-        image_info_list.append((dtb_offset, dtb_size, 'dtb'))
+        image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
 
-    if boot_signature_size > 0:
+    if info.boot_signature_size > 0:
         # boot signature only exists in boot.img version >= v4.
         # There are only kernel and ramdisk pages before the signature.
         boot_signature_offset = page_size * (
             num_header_pages + num_kernel_pages + num_ramdisk_pages)
 
-        image_info_list.append((boot_signature_offset, boot_signature_size,
+        image_info_list.append((boot_signature_offset, info.boot_signature_size,
                                 'boot_signature'))
 
-    for image_info in image_info_list:
-        extract_image(image_info[0], image_info[1], args.boot_img,
-                      os.path.join(args.out, image_info[2]))
+    create_out_dir(args.out)
+    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
+
+    return info
 
 
 class VendorBootImageInfoFormatter:
@@ -326,7 +335,7 @@
 
         return '\n'.join(lines)
 
-    def format_mkbootimg_argument(self, null=False):
+    def format_mkbootimg_argument(self):
         args = []
         args.extend(['--header_version', str(self.header_version)])
         args.extend(['--pagesize', f'{self.page_size:#010x}'])
@@ -338,12 +347,11 @@
         args.extend(['--vendor_cmdline', self.cmdline])
         args.extend(['--board', self.product_name])
 
-        dtb_path = os.path.join(self.image_dir, 'dtb')
-        args.extend(['--dtb', dtb_path])
+        args.extend(['--dtb', os.path.join(self.image_dir, 'dtb')])
 
         if self.header_version > 3:
-            bootconfig_path = os.path.join(self.image_dir, 'bootconfig')
-            args.extend(['--vendor_bootconfig', bootconfig_path])
+            args.extend(['--vendor_bootconfig',
+                         os.path.join(self.image_dir, 'bootconfig')])
 
             for entry in self.vendor_ramdisk_table:
                 (output_ramdisk_name, _, _, ramdisk_type,
@@ -357,38 +365,10 @@
                     self.image_dir, output_ramdisk_name)
                 args.extend(['--vendor_ramdisk_fragment', vendor_ramdisk_path])
         else:
-            vendor_ramdisk_path = os.path.join(self.image_dir, 'vendor_ramdisk')
-            args.extend(['--vendor_ramdisk', vendor_ramdisk_path])
+            args.extend(['--vendor_ramdisk',
+                         os.path.join(self.image_dir, 'vendor_ramdisk')])
 
-        if null:
-            return '\0'.join(args) + '\0'
-        return shlex.join(args)
-
-    def format_json_dict(self):
-        """Returns a dict of arguments to be used in mkbootimg.py later."""
-        args_dict = {}
-        args_dict['header_version'] = str(self.header_version)
-
-        # Format uint32_t as '0xFFFFFFFF', uint64_t as '0xFFFFFFFFEEEEEEEE'.
-        args_dict['pagesize'] = f'{self.page_size:#010x}'
-
-        # Kernel load address is base + kernel_offset in mkbootimg.py.
-        # However, we don't know the value of 'base' when unpacking a
-        # vendor_boot.img in this script. So always set 'base' to be zero and
-        # 'kernel_offset' to be the kernel load address. Same for
-        # 'ramdisk_offset', 'tags_offset' and 'dtb_offset'.
-        args_dict['base'] = f'{0:#010x}'
-        args_dict['kernel_offset'] = f'{self.kernel_load_address:#010x}'
-        args_dict['ramdisk_offset'] = f'{self.ramdisk_load_address:#010x}'
-        args_dict['tags_offset'] = f'{self.tags_load_address:#010x}'
-        # The type of dtb_offset is uint64_t.
-        args_dict['dtb_offset'] = f'{self.dtb_load_address:#018x}'
-
-        args_dict['vendor_cmdline'] = self.cmdline
-        args_dict['board'] = self.product_name
-
-        # TODO(bowgotsai): support for multiple vendor ramdisk (vendor boot v4).
-        return args_dict
+        return args
 
 
 def unpack_vendor_boot_image(args):
@@ -467,9 +447,10 @@
                              ) # header + vendor_ramdisk
     image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
 
-    for image_info in image_info_list:
-        extract_image(image_info[0], image_info[1], args.boot_img,
-                      os.path.join(args.out, image_info[2]))
+    create_out_dir(args.out)
+    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
 
     if info.header_version > 3:
         vendor_ramdisk_by_name_dir = os.path.join(
@@ -483,27 +464,27 @@
                 os.remove(dst_pathname)
             os.symlink(src_pathname, dst_pathname)
 
-    info.image_dir = args.out
-
-    # Saves the arguments to be reused in mkbootimg.py later.
-    mkbootimg_args = info.format_json_dict()
-    with open(os.path.join(args.out, MKBOOTIMG_ARGS_FILE), 'w') as f:
-        json.dump(mkbootimg_args, f, sort_keys=True, indent=4)
-
-    if args.format == 'mkbootimg':
-        print(info.format_mkbootimg_argument(null=args.null),
-              end='' if args.null else None)
-    else:
-        print(info.format_pretty_text())
+    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!':
-        unpack_boot_image(args)
+        info = unpack_boot_image(args)
     elif boot_magic == 'VNDRBOOT':
-        unpack_vendor_boot_image(args)
+        info = unpack_vendor_boot_image(args)
+    else:
+        raise ValueError(f'Not an Android boot image, magic: {boot_magic}')
+
+    if args.format == 'mkbootimg':
+        mkbootimg_args = info.format_mkbootimg_argument()
+        if args.null:
+            print('\0'.join(mkbootimg_args) + '\0', end='')
+        else:
+            print(shlex.join(mkbootimg_args))
+    else:
+        print(info.format_pretty_text())
 
 
 def get_unpack_usage():
@@ -559,7 +540,6 @@
 def main():
     """parse arguments and unpack boot image"""
     args = parse_cmdline()
-    create_out_dir(args.out)
     unpack_image(args)