Merge "repack_bootimg: changes the default path of --ramdisk_add" am: c1d56bbd1b am: bca59d92ab am: 9e3a1e8d4c am: b8021f3120 am: eacd6f2596

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

Change-Id: I28d5c1441d10492f8d7c3695734c4de9dede17e2
diff --git a/repack_bootimg.py b/repack_bootimg.py
index cd0732a..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,63 +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"""
 
@@ -237,16 +229,32 @@
         # 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))
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 0a508a9..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"""
@@ -130,58 +127,52 @@
 
         return '\n'.join(lines)
 
-    def _format_json_dict_boot_image_v2_and_below(self):
-        """Returns a dict of mkbootimg.py arguments for v0-v2 boot.img."""
-        args_dict = {}
+    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_dict['header_version'] = str(self.header_version)
-        # The type of pagesize is uint32_t, using '0xFFFFFFFF' as the
-        # output format.
-        args_dict['pagesize'] = f'{self.page_size:#010x}'
+        args.extend(['--kernel', os.path.join(self.image_dir, 'kernel')])
+        args.extend(['--ramdisk', os.path.join(self.image_dir, 'ramdisk')])
 
-        # 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', etc.
-        # The following types are uint32_t, using '0xFFFFFFFF' as the output
-        # format.
-        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['second_offset'] = f'{self.second_load_address:#010x}'
-        args_dict['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:
-            # The type of dtb_offset is uint64_t, using '0xFFFFFFFFEEEEEEEE' as
-            # the output format.
-            args_dict['dtb_offset'] = f'{self.dtb_load_address:#018x}'
-
-        args_dict['os_version'] = self.os_version
-        args_dict['os_patch_level'] = self.os_patch_level
-        args_dict['cmdline'] = self.cmdline + self.extra_cmdline
-        args_dict['board'] = self.product_name
-
-        return args_dict
-
-    def _format_json_dict_boot_image_v3_and_above(self):
-        """Returns a dict of mkbootimg.py arguments for >= v3 boot.img."""
-        args_dict = {}
-
-        args_dict['header_version'] = str(self.header_version)
-        args_dict['os_version'] = self.os_version
-        args_dict['os_patch_level'] = self.os_patch_level
-        args_dict['cmdline'] = self.cmdline
-
-        return args_dict
-
-    def format_json_dict(self):
-        """Returns a dict of arguments to be used in mkbootimg.py later."""
         if self.header_version <= 2:
-            return self._format_json_dict_boot_image_v2_and_below()
+            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:
-            return self._format_json_dict_boot_image_v3_and_above()
+            args.extend(['--cmdline', self.cmdline])
+
+        return args
 
 
 def unpack_boot_image(args):
@@ -287,17 +278,12 @@
         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
 
-    # 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)
-
-    # TODO(yochiang): Support --format=mkbootimg
-    print(info.format_pretty_text())
+    return info
 
 
 class VendorBootImageInfoFormatter:
@@ -349,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}'])
@@ -361,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,
@@ -380,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):
@@ -490,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(
@@ -506,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():
@@ -582,7 +540,6 @@
 def main():
     """parse arguments and unpack boot image"""
     args = parse_cmdline()
-    create_out_dir(args.out)
     unpack_image(args)