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