1#!/usr/bin/env python3
2#
3# Copyright 2020, The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Tests mkbootimg and unpack_bootimg."""
18
19import filecmp
20import logging
21import os
22import random
23import shlex
24import subprocess
25import sys
26import tempfile
27import unittest
28
29BOOT_ARGS_OFFSET = 64
30BOOT_ARGS_SIZE = 512
31BOOT_EXTRA_ARGS_OFFSET = 608
32BOOT_EXTRA_ARGS_SIZE = 1024
33BOOT_V3_ARGS_OFFSET = 44
34VENDOR_BOOT_ARGS_OFFSET = 28
35VENDOR_BOOT_ARGS_SIZE = 2048
36
37BOOT_IMAGE_V4_SIGNATURE_SIZE = 4096
38
39TEST_KERNEL_CMDLINE = (
40    'printk.devkmsg=on firmware_class.path=/vendor/etc/ init=/init '
41    'kfence.sample_interval=500 loop.max_part=7 bootconfig'
42)
43
44
45def generate_test_file(pathname, size, seed=None):
46    """Generates a gibberish-filled test file and returns its pathname."""
47    random.seed(os.path.basename(pathname) if seed is None else seed)
48    with open(pathname, 'wb') as f:
49        f.write(random.randbytes(size))
50    return pathname
51
52
53def subsequence_of(list1, list2):
54    """Returns True if list1 is a subsequence of list2.
55
56    >>> subsequence_of([], [1])
57    True
58    >>> subsequence_of([2, 4], [1, 2, 3, 4])
59    True
60    >>> subsequence_of([1, 2, 2], [1, 2, 3])
61    False
62    """
63    if len(list1) == 0:
64        return True
65    if len(list2) == 0:
66        return False
67    if list1[0] == list2[0]:
68        return subsequence_of(list1[1:], list2[1:])
69    return subsequence_of(list1, list2[1:])
70
71
72class MkbootimgTest(unittest.TestCase):
73    """Tests the functionalities of mkbootimg and unpack_bootimg."""
74
75    def setUp(self):
76        # Saves the test executable directory so that relative path references
77        # to test dependencies don't rely on being manually run from the
78        # executable directory.
79        # With this, we can just open "./tests/data/testkey_rsa2048.pem" in the
80        # following tests with subprocess.run(..., cwd=self._exec_dir, ...).
81        self._exec_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
82
83        self._avbtool_path = os.path.join(self._exec_dir, 'avbtool')
84
85        # Set self.maxDiff to None to see full diff in assertion.
86        # C0103: invalid-name for maxDiff.
87        self.maxDiff = None  # pylint: disable=C0103
88
89    def _test_boot_image_v4_signature(self, avbtool_path):
90        """Tests the boot_signature in boot.img v4."""
91        with tempfile.TemporaryDirectory() as temp_out_dir:
92            boot_img = os.path.join(temp_out_dir, 'boot.img')
93            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
94                                        0x1000)
95            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
96                                         0x1000)
97            mkbootimg_cmds = [
98                'mkbootimg',
99                '--header_version', '4',
100                '--kernel', kernel,
101                '--ramdisk', ramdisk,
102                '--cmdline', TEST_KERNEL_CMDLINE,
103                '--os_version', '11.0.0',
104                '--os_patch_level', '2021-01',
105                '--gki_signing_algorithm', 'SHA256_RSA2048',
106                '--gki_signing_key', './tests/data/testkey_rsa2048.pem',
107                '--gki_signing_signature_args',
108                '--prop foo:bar --prop gki:nice',
109                '--output', boot_img,
110            ]
111
112            if avbtool_path:
113                mkbootimg_cmds.extend(
114                    ['--gki_signing_avbtool_path', avbtool_path])
115
116            unpack_bootimg_cmds = [
117                'unpack_bootimg',
118                '--boot_img', boot_img,
119                '--out', os.path.join(temp_out_dir, 'out'),
120            ]
121
122            # cwd=self._exec_dir is required to read
123            # ./tests/data/testkey_rsa2048.pem for --gki_signing_key.
124            subprocess.run(mkbootimg_cmds, check=True, cwd=self._exec_dir)
125            subprocess.run(unpack_bootimg_cmds, check=True)
126
127            # Checks the content of the boot signature.
128            expected_boot_signature_info = (
129                'Minimum libavb version:   1.0\n'
130                'Header Block:             256 bytes\n'
131                'Authentication Block:     320 bytes\n'
132                'Auxiliary Block:          832 bytes\n'
133                'Public key (sha1):        '
134                'cdbb77177f731920bbe0a0f94f84d9038ae0617d\n'
135                'Algorithm:                SHA256_RSA2048\n'
136                'Rollback Index:           0\n'
137                'Flags:                    0\n'
138                'Rollback Index Location:  0\n'
139                "Release String:           'avbtool 1.2.0'\n"
140                'Descriptors:\n'
141                '    Hash descriptor:\n'
142                '      Image Size:            12288 bytes\n'
143                '      Hash Algorithm:        sha256\n'
144                '      Partition Name:        boot\n'
145                '      Salt:                  d00df00d\n'
146                '      Digest:                '
147                'cf3755630856f23ab70e501900050fee'
148                'f30b633b3e82a9085a578617e344f9c7\n'
149                '      Flags:                 0\n'
150                "    Prop: foo -> 'bar'\n"
151                "    Prop: gki -> 'nice'\n"
152            )
153
154            avbtool_info_cmds = [
155                # use avbtool_path if it is not None.
156                avbtool_path or 'avbtool',
157                'info_image', '--image',
158                os.path.join(temp_out_dir, 'out', 'boot_signature')
159            ]
160            result = subprocess.run(avbtool_info_cmds, check=True,
161                                    capture_output=True, encoding='utf-8')
162
163            self.assertEqual(result.stdout, expected_boot_signature_info)
164
165    def test_boot_image_v4_signature_without_avbtool_path(self):
166        """Boot signature generation without --gki_signing_avbtool_path."""
167        self._test_boot_image_v4_signature(avbtool_path=None)
168
169    def test_boot_image_v4_signature_with_avbtool_path(self):
170        """Boot signature generation with --gki_signing_avbtool_path."""
171        self._test_boot_image_v4_signature(avbtool_path=self._avbtool_path)
172
173    def test_boot_image_v4_signature_exceed_size(self):
174        """Tests the boot signature size exceeded in a boot image version 4."""
175        with tempfile.TemporaryDirectory() as temp_out_dir:
176            boot_img = os.path.join(temp_out_dir, 'boot.img')
177            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
178                                        0x1000)
179            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
180                                         0x1000)
181            mkbootimg_cmds = [
182                'mkbootimg',
183                '--header_version', '4',
184                '--kernel', kernel,
185                '--ramdisk', ramdisk,
186                '--cmdline', TEST_KERNEL_CMDLINE,
187                '--os_version', '11.0.0',
188                '--os_patch_level', '2021-01',
189                '--gki_signing_avbtool_path', self._avbtool_path,
190                '--gki_signing_algorithm', 'SHA256_RSA2048',
191                '--gki_signing_key', './tests/data/testkey_rsa2048.pem',
192                '--gki_signing_signature_args',
193                # Makes it exceed the signature max size.
194                '--prop foo:bar --prop gki:nice ' * 64,
195                '--output', boot_img,
196            ]
197
198            # cwd=self._exec_dir is required to read
199            # ./tests/data/testkey_rsa2048.pem for --gki_signing_key.
200            try:
201                subprocess.run(mkbootimg_cmds, check=True, capture_output=True,
202                               cwd=self._exec_dir, encoding='utf-8')
203                self.fail('Exceeding signature size assertion is not raised')
204            except subprocess.CalledProcessError as e:
205                self.assertIn('ValueError: boot sigature size is > 4096',
206                              e.stderr)
207
208    def test_boot_image_v4_signature_zeros(self):
209        """Tests no boot signature in a boot image version 4."""
210        with tempfile.TemporaryDirectory() as temp_out_dir:
211            boot_img = os.path.join(temp_out_dir, 'boot.img')
212            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
213                                        0x1000)
214            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
215                                         0x1000)
216
217            # The boot signature will be zeros if no
218            # --gki_signing_[algorithm|key] is provided.
219            mkbootimg_cmds = [
220                'mkbootimg',
221                '--header_version', '4',
222                '--kernel', kernel,
223                '--ramdisk', ramdisk,
224                '--cmdline', TEST_KERNEL_CMDLINE,
225                '--os_version', '11.0.0',
226                '--os_patch_level', '2021-01',
227                '--output', boot_img,
228            ]
229            unpack_bootimg_cmds = [
230                'unpack_bootimg',
231                '--boot_img', boot_img,
232                '--out', os.path.join(temp_out_dir, 'out'),
233            ]
234
235            subprocess.run(mkbootimg_cmds, check=True)
236            subprocess.run(unpack_bootimg_cmds, check=True)
237
238            boot_signature = os.path.join(
239                temp_out_dir, 'out', 'boot_signature')
240            with open(boot_signature) as f:
241                zeros = '\x00' * BOOT_IMAGE_V4_SIGNATURE_SIZE
242                self.assertEqual(f.read(), zeros)
243
244    def test_vendor_boot_v4(self):
245        """Tests vendor_boot version 4."""
246        with tempfile.TemporaryDirectory() as temp_out_dir:
247            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
248            dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
249            ramdisk1 = generate_test_file(
250                os.path.join(temp_out_dir, 'ramdisk1'), 0x1000)
251            ramdisk2 = generate_test_file(
252                os.path.join(temp_out_dir, 'ramdisk2'), 0x2000)
253            bootconfig = generate_test_file(
254                os.path.join(temp_out_dir, 'bootconfig'), 0x1000)
255            mkbootimg_cmds = [
256                'mkbootimg',
257                '--header_version', '4',
258                '--vendor_boot', vendor_boot_img,
259                '--dtb', dtb,
260                '--vendor_ramdisk', ramdisk1,
261                '--ramdisk_type', 'PLATFORM',
262                '--ramdisk_name', 'RAMDISK1',
263                '--vendor_ramdisk_fragment', ramdisk1,
264                '--ramdisk_type', 'DLKM',
265                '--ramdisk_name', 'RAMDISK2',
266                '--board_id0', '0xC0FFEE',
267                '--board_id15', '0x15151515',
268                '--vendor_ramdisk_fragment', ramdisk2,
269                '--vendor_cmdline', TEST_KERNEL_CMDLINE,
270                '--vendor_bootconfig', bootconfig,
271            ]
272            unpack_bootimg_cmds = [
273                'unpack_bootimg',
274                '--boot_img', vendor_boot_img,
275                '--out', os.path.join(temp_out_dir, 'out'),
276            ]
277            expected_output = [
278                'boot magic: VNDRBOOT',
279                'vendor boot image header version: 4',
280                'vendor ramdisk total size: 16384',
281                f'vendor command line args: {TEST_KERNEL_CMDLINE}',
282                'dtb size: 4096',
283                'vendor ramdisk table size: 324',
284                'size: 4096', 'offset: 0', 'type: 0x1', 'name:',
285                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
286                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
287                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
288                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
289                'size: 4096', 'offset: 4096', 'type: 0x1', 'name: RAMDISK1',
290                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
291                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
292                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
293                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
294                'size: 8192', 'offset: 8192', 'type: 0x3', 'name: RAMDISK2',
295                '0x00c0ffee, 0x00000000, 0x00000000, 0x00000000,',
296                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
297                '0x00000000, 0x00000000, 0x00000000, 0x00000000,',
298                '0x00000000, 0x00000000, 0x00000000, 0x15151515,',
299                'vendor bootconfig size: 4096',
300            ]
301
302            subprocess.run(mkbootimg_cmds, check=True)
303            result = subprocess.run(unpack_bootimg_cmds, check=True,
304                                    capture_output=True, encoding='utf-8')
305            output = [line.strip() for line in result.stdout.splitlines()]
306            if not subsequence_of(expected_output, output):
307                msg = '\n'.join([
308                    'Unexpected unpack_bootimg output:',
309                    'Expected:',
310                    ' ' + '\n '.join(expected_output),
311                    '',
312                    'Actual:',
313                    ' ' + '\n '.join(output),
314                ])
315                self.fail(msg)
316
317    def test_unpack_vendor_boot_image_v4(self):
318        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
319        with tempfile.TemporaryDirectory() as temp_out_dir:
320            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
321            vendor_boot_img_reconstructed = os.path.join(
322                temp_out_dir, 'vendor_boot.img.reconstructed')
323            dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
324            ramdisk1 = generate_test_file(
325                os.path.join(temp_out_dir, 'ramdisk1'), 0x121212)
326            ramdisk2 = generate_test_file(
327                os.path.join(temp_out_dir, 'ramdisk2'), 0x212121)
328            bootconfig = generate_test_file(
329                os.path.join(temp_out_dir, 'bootconfig'), 0x1000)
330
331            mkbootimg_cmds = [
332                'mkbootimg',
333                '--header_version', '4',
334                '--vendor_boot', vendor_boot_img,
335                '--dtb', dtb,
336                '--vendor_ramdisk', ramdisk1,
337                '--ramdisk_type', 'PLATFORM',
338                '--ramdisk_name', 'RAMDISK1',
339                '--vendor_ramdisk_fragment', ramdisk1,
340                '--ramdisk_type', 'DLKM',
341                '--ramdisk_name', 'RAMDISK2',
342                '--board_id0', '0xC0FFEE',
343                '--board_id15', '0x15151515',
344                '--vendor_ramdisk_fragment', ramdisk2,
345                '--vendor_cmdline', TEST_KERNEL_CMDLINE,
346                '--vendor_bootconfig', bootconfig,
347            ]
348            unpack_bootimg_cmds = [
349                'unpack_bootimg',
350                '--boot_img', vendor_boot_img,
351                '--out', os.path.join(temp_out_dir, 'out'),
352                '--format=mkbootimg',
353            ]
354            subprocess.run(mkbootimg_cmds, check=True)
355            result = subprocess.run(unpack_bootimg_cmds, check=True,
356                                    capture_output=True, encoding='utf-8')
357            mkbootimg_cmds = [
358                'mkbootimg',
359                '--vendor_boot', vendor_boot_img_reconstructed,
360            ]
361            unpack_format_args = shlex.split(result.stdout)
362            mkbootimg_cmds.extend(unpack_format_args)
363
364            subprocess.run(mkbootimg_cmds, check=True)
365            self.assertTrue(
366                filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
367                'reconstructed vendor_boot image differ from the original')
368
369            # Also check that -0, --null are as expected.
370            unpack_bootimg_cmds.append('--null')
371            result = subprocess.run(unpack_bootimg_cmds, check=True,
372                                    capture_output=True, encoding='utf-8')
373            unpack_format_null_args = result.stdout
374            self.assertEqual('\0'.join(unpack_format_args) + '\0',
375                             unpack_format_null_args)
376
377    def test_unpack_vendor_boot_image_v3(self):
378        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
379        with tempfile.TemporaryDirectory() as temp_out_dir:
380            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
381            vendor_boot_img_reconstructed = os.path.join(
382                temp_out_dir, 'vendor_boot.img.reconstructed')
383            dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
384            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
385                                         0x121212)
386            mkbootimg_cmds = [
387                'mkbootimg',
388                '--header_version', '3',
389                '--vendor_boot', vendor_boot_img,
390                '--vendor_ramdisk', ramdisk,
391                '--dtb', dtb,
392                '--vendor_cmdline', TEST_KERNEL_CMDLINE,
393                '--board', 'product_name',
394                '--base', '0x00000000',
395                '--dtb_offset', '0x01f00000',
396                '--kernel_offset', '0x00008000',
397                '--pagesize', '0x00001000',
398                '--ramdisk_offset', '0x01000000',
399                '--tags_offset', '0x00000100',
400            ]
401            unpack_bootimg_cmds = [
402                'unpack_bootimg',
403                '--boot_img', vendor_boot_img,
404                '--out', os.path.join(temp_out_dir, 'out'),
405                '--format=mkbootimg',
406            ]
407            subprocess.run(mkbootimg_cmds, check=True)
408            result = subprocess.run(unpack_bootimg_cmds, check=True,
409                                    capture_output=True, encoding='utf-8')
410            mkbootimg_cmds = [
411                'mkbootimg',
412                '--vendor_boot', vendor_boot_img_reconstructed,
413            ]
414            mkbootimg_cmds.extend(shlex.split(result.stdout))
415
416            subprocess.run(mkbootimg_cmds, check=True)
417            self.assertTrue(
418                filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
419                'reconstructed vendor_boot image differ from the original')
420
421    def test_unpack_boot_image_v3(self):
422        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
423        with tempfile.TemporaryDirectory() as temp_out_dir:
424            boot_img = os.path.join(temp_out_dir, 'boot.img')
425            boot_img_reconstructed = os.path.join(
426                temp_out_dir, 'boot.img.reconstructed')
427            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
428                                        0x1000)
429            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
430                                         0x1000)
431            mkbootimg_cmds = [
432                'mkbootimg',
433                '--header_version', '3',
434                '--kernel', kernel,
435                '--ramdisk', ramdisk,
436                '--cmdline', TEST_KERNEL_CMDLINE,
437                '--os_version', '11.0.0',
438                '--os_patch_level', '2021-01',
439                '--output', boot_img,
440            ]
441            unpack_bootimg_cmds = [
442                'unpack_bootimg',
443                '--boot_img', boot_img,
444                '--out', os.path.join(temp_out_dir, 'out'),
445                '--format=mkbootimg',
446            ]
447
448            subprocess.run(mkbootimg_cmds, check=True)
449            result = subprocess.run(unpack_bootimg_cmds, check=True,
450                                    capture_output=True, encoding='utf-8')
451            mkbootimg_cmds = [
452                'mkbootimg',
453                '--out', boot_img_reconstructed,
454            ]
455            mkbootimg_cmds.extend(shlex.split(result.stdout))
456
457            subprocess.run(mkbootimg_cmds, check=True)
458            self.assertTrue(
459                filecmp.cmp(boot_img, boot_img_reconstructed),
460                'reconstructed boot image differ from the original')
461
462    def test_unpack_boot_image_v2(self):
463        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
464        with tempfile.TemporaryDirectory() as temp_out_dir:
465            # Output image path.
466            boot_img = os.path.join(temp_out_dir, 'boot.img')
467            boot_img_reconstructed = os.path.join(
468                temp_out_dir, 'boot.img.reconstructed')
469            # Creates blank images first.
470            kernel = generate_test_file(
471                os.path.join(temp_out_dir, 'kernel'), 0x1000)
472            ramdisk = generate_test_file(
473                os.path.join(temp_out_dir, 'ramdisk'), 0x1000)
474            second = generate_test_file(
475                os.path.join(temp_out_dir, 'second'), 0x1000)
476            recovery_dtbo = generate_test_file(
477                os.path.join(temp_out_dir, 'recovery_dtbo'), 0x1000)
478            dtb = generate_test_file(
479                os.path.join(temp_out_dir, 'dtb'), 0x1000)
480
481            cmdline = (BOOT_ARGS_SIZE - 1) * 'x'
482            extra_cmdline = (BOOT_EXTRA_ARGS_SIZE - 1) * 'y'
483
484            mkbootimg_cmds = [
485                'mkbootimg',
486                '--header_version', '2',
487                '--base', '0x00000000',
488                '--kernel', kernel,
489                '--kernel_offset', '0x00008000',
490                '--ramdisk', ramdisk,
491                '--ramdisk_offset', '0x01000000',
492                '--second', second,
493                '--second_offset', '0x40000000',
494                '--recovery_dtbo', recovery_dtbo,
495                '--dtb', dtb,
496                '--dtb_offset', '0x01f00000',
497                '--tags_offset', '0x00000100',
498                '--pagesize', '0x00001000',
499                '--os_version', '11.0.0',
500                '--os_patch_level', '2021-03',
501                '--board', 'boot_v2',
502                '--cmdline', cmdline + extra_cmdline,
503                '--output', boot_img,
504            ]
505            unpack_bootimg_cmds = [
506                'unpack_bootimg',
507                '--boot_img', boot_img,
508                '--out', os.path.join(temp_out_dir, 'out'),
509                '--format=mkbootimg',
510            ]
511
512            subprocess.run(mkbootimg_cmds, check=True)
513            result = subprocess.run(unpack_bootimg_cmds, check=True,
514                                    capture_output=True, encoding='utf-8')
515            mkbootimg_cmds = [
516                'mkbootimg',
517                '--out', boot_img_reconstructed,
518            ]
519            mkbootimg_cmds.extend(shlex.split(result.stdout))
520
521            subprocess.run(mkbootimg_cmds, check=True)
522            self.assertTrue(
523                filecmp.cmp(boot_img, boot_img_reconstructed),
524                'reconstructed boot image differ from the original')
525
526    def test_unpack_boot_image_v1(self):
527        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
528        with tempfile.TemporaryDirectory() as temp_out_dir:
529            # Output image path.
530            boot_img = os.path.join(temp_out_dir, 'boot.img')
531            boot_img_reconstructed = os.path.join(
532                temp_out_dir, 'boot.img.reconstructed')
533            # Creates blank images first.
534            kernel = generate_test_file(
535                os.path.join(temp_out_dir, 'kernel'), 0x1000)
536            ramdisk = generate_test_file(
537                os.path.join(temp_out_dir, 'ramdisk'), 0x1000)
538            recovery_dtbo = generate_test_file(
539                os.path.join(temp_out_dir, 'recovery_dtbo'), 0x1000)
540
541            cmdline = (BOOT_ARGS_SIZE - 1) * 'x'
542            extra_cmdline = (BOOT_EXTRA_ARGS_SIZE - 1) * 'y'
543
544            mkbootimg_cmds = [
545                'mkbootimg',
546                '--header_version', '1',
547                '--base', '0x00000000',
548                '--kernel', kernel,
549                '--kernel_offset', '0x00008000',
550                '--ramdisk', ramdisk,
551                '--ramdisk_offset', '0x01000000',
552                '--recovery_dtbo', recovery_dtbo,
553                '--tags_offset', '0x00000100',
554                '--pagesize', '0x00001000',
555                '--os_version', '11.0.0',
556                '--os_patch_level', '2021-03',
557                '--board', 'boot_v1',
558                '--cmdline', cmdline + extra_cmdline,
559                '--output', boot_img,
560            ]
561            unpack_bootimg_cmds = [
562                'unpack_bootimg',
563                '--boot_img', boot_img,
564                '--out', os.path.join(temp_out_dir, 'out'),
565                '--format=mkbootimg',
566            ]
567
568            subprocess.run(mkbootimg_cmds, check=True)
569            result = subprocess.run(unpack_bootimg_cmds, check=True,
570                                    capture_output=True, encoding='utf-8')
571            mkbootimg_cmds = [
572                'mkbootimg',
573                '--out', boot_img_reconstructed,
574            ]
575            mkbootimg_cmds.extend(shlex.split(result.stdout))
576
577            subprocess.run(mkbootimg_cmds, check=True)
578            self.assertTrue(
579                filecmp.cmp(boot_img, boot_img_reconstructed),
580                'reconstructed boot image differ from the original')
581
582    def test_unpack_boot_image_v0(self):
583        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
584        with tempfile.TemporaryDirectory() as temp_out_dir:
585            # Output image path.
586            boot_img = os.path.join(temp_out_dir, 'boot.img')
587            boot_img_reconstructed = os.path.join(
588                temp_out_dir, 'boot.img.reconstructed')
589            # Creates blank images first.
590            kernel = generate_test_file(
591                os.path.join(temp_out_dir, 'kernel'), 0x1000)
592            ramdisk = generate_test_file(
593                os.path.join(temp_out_dir, 'ramdisk'), 0x1000)
594            second = generate_test_file(
595                os.path.join(temp_out_dir, 'second'), 0x1000)
596
597            cmdline = (BOOT_ARGS_SIZE - 1) * 'x'
598            extra_cmdline = (BOOT_EXTRA_ARGS_SIZE - 1) * 'y'
599
600            mkbootimg_cmds = [
601                'mkbootimg',
602                '--header_version', '0',
603                '--base', '0x00000000',
604                '--kernel', kernel,
605                '--kernel_offset', '0x00008000',
606                '--ramdisk', ramdisk,
607                '--ramdisk_offset', '0x01000000',
608                '--second', second,
609                '--second_offset', '0x40000000',
610                '--tags_offset', '0x00000100',
611                '--pagesize', '0x00001000',
612                '--os_version', '11.0.0',
613                '--os_patch_level', '2021-03',
614                '--board', 'boot_v0',
615                '--cmdline', cmdline + extra_cmdline,
616                '--output', boot_img,
617            ]
618            unpack_bootimg_cmds = [
619                'unpack_bootimg',
620                '--boot_img', boot_img,
621                '--out', os.path.join(temp_out_dir, 'out'),
622            ]
623            unpack_bootimg_cmds = [
624                'unpack_bootimg',
625                '--boot_img', boot_img,
626                '--out', os.path.join(temp_out_dir, 'out'),
627                '--format=mkbootimg',
628            ]
629
630            subprocess.run(mkbootimg_cmds, check=True)
631            result = subprocess.run(unpack_bootimg_cmds, check=True,
632                                    capture_output=True, encoding='utf-8')
633            mkbootimg_cmds = [
634                'mkbootimg',
635                '--out', boot_img_reconstructed,
636            ]
637            mkbootimg_cmds.extend(shlex.split(result.stdout))
638
639            subprocess.run(mkbootimg_cmds, check=True)
640            self.assertTrue(
641                filecmp.cmp(boot_img, boot_img_reconstructed),
642                'reconstructed boot image differ from the original')
643
644    def test_boot_image_v2_cmdline_null_terminator(self):
645        """Tests that kernel commandline is null-terminated."""
646        with tempfile.TemporaryDirectory() as temp_out_dir:
647            dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
648            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
649                                        0x1000)
650            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
651                                         0x1000)
652            cmdline = (BOOT_ARGS_SIZE - 1) * 'x'
653            extra_cmdline = (BOOT_EXTRA_ARGS_SIZE - 1) * 'y'
654            boot_img = os.path.join(temp_out_dir, 'boot.img')
655            mkbootimg_cmds = [
656                'mkbootimg',
657                '--header_version', '2',
658                '--dtb', dtb,
659                '--kernel', kernel,
660                '--ramdisk', ramdisk,
661                '--cmdline', cmdline + extra_cmdline,
662                '--output', boot_img,
663            ]
664
665            subprocess.run(mkbootimg_cmds, check=True)
666
667            with open(boot_img, 'rb') as f:
668                raw_boot_img = f.read()
669            raw_cmdline = raw_boot_img[BOOT_ARGS_OFFSET:][:BOOT_ARGS_SIZE]
670            raw_extra_cmdline = (raw_boot_img[BOOT_EXTRA_ARGS_OFFSET:]
671                                 [:BOOT_EXTRA_ARGS_SIZE])
672            self.assertEqual(raw_cmdline, cmdline.encode() + b'\x00')
673            self.assertEqual(raw_extra_cmdline,
674                             extra_cmdline.encode() + b'\x00')
675
676    def test_boot_image_v3_cmdline_null_terminator(self):
677        """Tests that kernel commandline is null-terminated."""
678        with tempfile.TemporaryDirectory() as temp_out_dir:
679            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
680                                        0x1000)
681            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
682                                         0x1000)
683            cmdline = BOOT_ARGS_SIZE * 'x' + (BOOT_EXTRA_ARGS_SIZE - 1) * 'y'
684            boot_img = os.path.join(temp_out_dir, 'boot.img')
685            mkbootimg_cmds = [
686                'mkbootimg',
687                '--header_version', '3',
688                '--kernel', kernel,
689                '--ramdisk', ramdisk,
690                '--cmdline', cmdline,
691                '--output', boot_img,
692            ]
693
694            subprocess.run(mkbootimg_cmds, check=True)
695
696            with open(boot_img, 'rb') as f:
697                raw_boot_img = f.read()
698            raw_cmdline = (raw_boot_img[BOOT_V3_ARGS_OFFSET:]
699                           [:BOOT_ARGS_SIZE + BOOT_EXTRA_ARGS_SIZE])
700            self.assertEqual(raw_cmdline, cmdline.encode() + b'\x00')
701
702    def test_vendor_boot_image_v3_cmdline_null_terminator(self):
703        """Tests that kernel commandline is null-terminated."""
704        with tempfile.TemporaryDirectory() as temp_out_dir:
705            dtb = generate_test_file(os.path.join(temp_out_dir, 'dtb'), 0x1000)
706            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
707                                         0x1000)
708            vendor_cmdline = (VENDOR_BOOT_ARGS_SIZE - 1) * 'x'
709            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
710            mkbootimg_cmds = [
711                'mkbootimg',
712                '--header_version', '3',
713                '--dtb', dtb,
714                '--vendor_ramdisk', ramdisk,
715                '--vendor_cmdline', vendor_cmdline,
716                '--vendor_boot', vendor_boot_img,
717            ]
718
719            subprocess.run(mkbootimg_cmds, check=True)
720
721            with open(vendor_boot_img, 'rb') as f:
722                raw_vendor_boot_img = f.read()
723            raw_vendor_cmdline = (raw_vendor_boot_img[VENDOR_BOOT_ARGS_OFFSET:]
724                                  [:VENDOR_BOOT_ARGS_SIZE])
725            self.assertEqual(raw_vendor_cmdline,
726                             vendor_cmdline.encode() + b'\x00')
727
728
729# I don't know how, but we need both the logger configuration and verbosity
730# level > 2 to make atest work. And yes this line needs to be at the very top
731# level, not even in the "__main__" indentation block.
732logging.basicConfig(stream=sys.stdout)
733
734if __name__ == '__main__':
735    unittest.main(verbosity=2)
736