1#!/usr/bin/env python3
2#
3# Copyright 2015, 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"""Creates the boot image."""
18
19from argparse import (ArgumentParser, ArgumentTypeError,
20                      FileType, RawDescriptionHelpFormatter)
21from hashlib import sha1
22from os import fstat
23from struct import pack
24
25import array
26import collections
27import os
28import re
29import tempfile
30
31from gki.generate_gki_certificate import generate_gki_certificate
32
33# Constant and structure definition is in
34# system/tools/mkbootimg/include/bootimg/bootimg.h
35BOOT_MAGIC = 'ANDROID!'
36BOOT_MAGIC_SIZE = 8
37BOOT_NAME_SIZE = 16
38BOOT_ARGS_SIZE = 512
39BOOT_EXTRA_ARGS_SIZE = 1024
40BOOT_IMAGE_HEADER_V1_SIZE = 1648
41BOOT_IMAGE_HEADER_V2_SIZE = 1660
42BOOT_IMAGE_HEADER_V3_SIZE = 1580
43BOOT_IMAGE_HEADER_V3_PAGESIZE = 4096
44BOOT_IMAGE_HEADER_V4_SIZE = 1584
45BOOT_IMAGE_V4_SIGNATURE_SIZE = 4096
46
47VENDOR_BOOT_MAGIC = 'VNDRBOOT'
48VENDOR_BOOT_MAGIC_SIZE = 8
49VENDOR_BOOT_NAME_SIZE = BOOT_NAME_SIZE
50VENDOR_BOOT_ARGS_SIZE = 2048
51VENDOR_BOOT_IMAGE_HEADER_V3_SIZE = 2112
52VENDOR_BOOT_IMAGE_HEADER_V4_SIZE = 2128
53
54VENDOR_RAMDISK_TYPE_NONE = 0
55VENDOR_RAMDISK_TYPE_PLATFORM = 1
56VENDOR_RAMDISK_TYPE_RECOVERY = 2
57VENDOR_RAMDISK_TYPE_DLKM = 3
58VENDOR_RAMDISK_NAME_SIZE = 32
59VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE = 16
60VENDOR_RAMDISK_TABLE_ENTRY_V4_SIZE = 108
61
62# Names with special meaning, mustn't be specified in --ramdisk_name.
63VENDOR_RAMDISK_NAME_BLOCKLIST = {b'default'}
64
65PARSER_ARGUMENT_VENDOR_RAMDISK_FRAGMENT = '--vendor_ramdisk_fragment'
66
67
68def filesize(f):
69    if f is None:
70        return 0
71    try:
72        return fstat(f.fileno()).st_size
73    except OSError:
74        return 0
75
76
77def update_sha(sha, f):
78    if f:
79        sha.update(f.read())
80        f.seek(0)
81        sha.update(pack('I', filesize(f)))
82    else:
83        sha.update(pack('I', 0))
84
85
86def pad_file(f, padding):
87    pad = (padding - (f.tell() & (padding - 1))) & (padding - 1)
88    f.write(pack(str(pad) + 'x'))
89
90
91def get_number_of_pages(image_size, page_size):
92    """calculates the number of pages required for the image"""
93    return (image_size + page_size - 1) // page_size
94
95
96def get_recovery_dtbo_offset(args):
97    """calculates the offset of recovery_dtbo image in the boot image"""
98    num_header_pages = 1 # header occupies a page
99    num_kernel_pages = get_number_of_pages(filesize(args.kernel), args.pagesize)
100    num_ramdisk_pages = get_number_of_pages(filesize(args.ramdisk),
101                                            args.pagesize)
102    num_second_pages = get_number_of_pages(filesize(args.second), args.pagesize)
103    dtbo_offset = args.pagesize * (num_header_pages + num_kernel_pages +
104                                   num_ramdisk_pages + num_second_pages)
105    return dtbo_offset
106
107
108def should_add_legacy_gki_boot_signature(args):
109    if args.gki_signing_key and args.gki_signing_algorithm:
110        return True
111    return False
112
113
114def write_header_v3_and_above(args):
115    if args.header_version > 3:
116        boot_header_size = BOOT_IMAGE_HEADER_V4_SIZE
117    else:
118        boot_header_size = BOOT_IMAGE_HEADER_V3_SIZE
119
120    args.output.write(pack(f'{BOOT_MAGIC_SIZE}s', BOOT_MAGIC.encode()))
121    # kernel size in bytes
122    args.output.write(pack('I', filesize(args.kernel)))
123    # ramdisk size in bytes
124    args.output.write(pack('I', filesize(args.ramdisk)))
125    # os version and patch level
126    args.output.write(pack('I', (args.os_version << 11) | args.os_patch_level))
127    args.output.write(pack('I', boot_header_size))
128    # reserved
129    args.output.write(pack('4I', 0, 0, 0, 0))
130    # version of boot image header
131    args.output.write(pack('I', args.header_version))
132    args.output.write(pack(f'{BOOT_ARGS_SIZE + BOOT_EXTRA_ARGS_SIZE}s',
133                           args.cmdline))
134    if args.header_version >= 4:
135        # The signature used to verify boot image v4.
136        boot_signature_size = 0
137        if should_add_legacy_gki_boot_signature(args):
138            boot_signature_size = BOOT_IMAGE_V4_SIGNATURE_SIZE
139        args.output.write(pack('I', boot_signature_size))
140    pad_file(args.output, BOOT_IMAGE_HEADER_V3_PAGESIZE)
141
142
143def write_vendor_boot_header(args):
144    if args.header_version > 3:
145        vendor_ramdisk_size = args.vendor_ramdisk_total_size
146        vendor_boot_header_size = VENDOR_BOOT_IMAGE_HEADER_V4_SIZE
147    else:
148        vendor_ramdisk_size = filesize(args.vendor_ramdisk)
149        vendor_boot_header_size = VENDOR_BOOT_IMAGE_HEADER_V3_SIZE
150
151    args.vendor_boot.write(pack(f'{VENDOR_BOOT_MAGIC_SIZE}s',
152                                VENDOR_BOOT_MAGIC.encode()))
153    # version of boot image header
154    args.vendor_boot.write(pack('I', args.header_version))
155    # flash page size
156    args.vendor_boot.write(pack('I', args.pagesize))
157    # kernel physical load address
158    args.vendor_boot.write(pack('I', args.base + args.kernel_offset))
159    # ramdisk physical load address
160    args.vendor_boot.write(pack('I', args.base + args.ramdisk_offset))
161    # ramdisk size in bytes
162    args.vendor_boot.write(pack('I', vendor_ramdisk_size))
163    args.vendor_boot.write(pack(f'{VENDOR_BOOT_ARGS_SIZE}s',
164                                args.vendor_cmdline))
165    # kernel tags physical load address
166    args.vendor_boot.write(pack('I', args.base + args.tags_offset))
167    # asciiz product name
168    args.vendor_boot.write(pack(f'{VENDOR_BOOT_NAME_SIZE}s', args.board))
169
170    # header size in bytes
171    args.vendor_boot.write(pack('I', vendor_boot_header_size))
172
173    # dtb size in bytes
174    args.vendor_boot.write(pack('I', filesize(args.dtb)))
175    # dtb physical load address
176    args.vendor_boot.write(pack('Q', args.base + args.dtb_offset))
177
178    if args.header_version > 3:
179        vendor_ramdisk_table_size = (args.vendor_ramdisk_table_entry_num *
180                                     VENDOR_RAMDISK_TABLE_ENTRY_V4_SIZE)
181        # vendor ramdisk table size in bytes
182        args.vendor_boot.write(pack('I', vendor_ramdisk_table_size))
183        # number of vendor ramdisk table entries
184        args.vendor_boot.write(pack('I', args.vendor_ramdisk_table_entry_num))
185        # vendor ramdisk table entry size in bytes
186        args.vendor_boot.write(pack('I', VENDOR_RAMDISK_TABLE_ENTRY_V4_SIZE))
187        # bootconfig section size in bytes
188        args.vendor_boot.write(pack('I', filesize(args.vendor_bootconfig)))
189    pad_file(args.vendor_boot, args.pagesize)
190
191
192def write_header(args):
193    if args.header_version > 4:
194        raise ValueError(
195            f'Boot header version {args.header_version} not supported')
196    if args.header_version in {3, 4}:
197        return write_header_v3_and_above(args)
198
199    ramdisk_load_address = ((args.base + args.ramdisk_offset)
200                            if filesize(args.ramdisk) > 0 else 0)
201    second_load_address = ((args.base + args.second_offset)
202                           if filesize(args.second) > 0 else 0)
203
204    args.output.write(pack(f'{BOOT_MAGIC_SIZE}s', BOOT_MAGIC.encode()))
205    # kernel size in bytes
206    args.output.write(pack('I', filesize(args.kernel)))
207    # kernel physical load address
208    args.output.write(pack('I', args.base + args.kernel_offset))
209    # ramdisk size in bytes
210    args.output.write(pack('I', filesize(args.ramdisk)))
211    # ramdisk physical load address
212    args.output.write(pack('I', ramdisk_load_address))
213    # second bootloader size in bytes
214    args.output.write(pack('I', filesize(args.second)))
215    # second bootloader physical load address
216    args.output.write(pack('I', second_load_address))
217    # kernel tags physical load address
218    args.output.write(pack('I', args.base + args.tags_offset))
219    # flash page size
220    args.output.write(pack('I', args.pagesize))
221    # version of boot image header
222    args.output.write(pack('I', args.header_version))
223    # os version and patch level
224    args.output.write(pack('I', (args.os_version << 11) | args.os_patch_level))
225    # asciiz product name
226    args.output.write(pack(f'{BOOT_NAME_SIZE}s', args.board))
227    args.output.write(pack(f'{BOOT_ARGS_SIZE}s', args.cmdline))
228
229    sha = sha1()
230    update_sha(sha, args.kernel)
231    update_sha(sha, args.ramdisk)
232    update_sha(sha, args.second)
233
234    if args.header_version > 0:
235        update_sha(sha, args.recovery_dtbo)
236    if args.header_version > 1:
237        update_sha(sha, args.dtb)
238
239    img_id = pack('32s', sha.digest())
240
241    args.output.write(img_id)
242    args.output.write(pack(f'{BOOT_EXTRA_ARGS_SIZE}s', args.extra_cmdline))
243
244    if args.header_version > 0:
245        if args.recovery_dtbo:
246            # recovery dtbo size in bytes
247            args.output.write(pack('I', filesize(args.recovery_dtbo)))
248            # recovert dtbo offset in the boot image
249            args.output.write(pack('Q', get_recovery_dtbo_offset(args)))
250        else:
251            # Set to zero if no recovery dtbo
252            args.output.write(pack('I', 0))
253            args.output.write(pack('Q', 0))
254
255    # Populate boot image header size for header versions 1 and 2.
256    if args.header_version == 1:
257        args.output.write(pack('I', BOOT_IMAGE_HEADER_V1_SIZE))
258    elif args.header_version == 2:
259        args.output.write(pack('I', BOOT_IMAGE_HEADER_V2_SIZE))
260
261    if args.header_version > 1:
262        if filesize(args.dtb) == 0:
263            raise ValueError('DTB image must not be empty.')
264
265        # dtb size in bytes
266        args.output.write(pack('I', filesize(args.dtb)))
267        # dtb physical load address
268        args.output.write(pack('Q', args.base + args.dtb_offset))
269
270    pad_file(args.output, args.pagesize)
271    return img_id
272
273
274class AsciizBytes:
275    """Parses a string and encodes it as an asciiz bytes object.
276
277    >>> AsciizBytes(bufsize=4)('foo')
278    b'foo\\x00'
279    >>> AsciizBytes(bufsize=4)('foob')
280    Traceback (most recent call last):
281        ...
282    argparse.ArgumentTypeError: Encoded asciiz length exceeded: max 4, got 5
283    """
284
285    def __init__(self, bufsize):
286        self.bufsize = bufsize
287
288    def __call__(self, arg):
289        arg_bytes = arg.encode() + b'\x00'
290        if len(arg_bytes) > self.bufsize:
291            raise ArgumentTypeError(
292                'Encoded asciiz length exceeded: '
293                f'max {self.bufsize}, got {len(arg_bytes)}')
294        return arg_bytes
295
296
297class VendorRamdiskTableBuilder:
298    """Vendor ramdisk table builder.
299
300    Attributes:
301        entries: A list of VendorRamdiskTableEntry namedtuple.
302        ramdisk_total_size: Total size in bytes of all ramdisks in the table.
303    """
304
305    VendorRamdiskTableEntry = collections.namedtuple(  # pylint: disable=invalid-name
306        'VendorRamdiskTableEntry',
307        ['ramdisk_path', 'ramdisk_size', 'ramdisk_offset', 'ramdisk_type',
308         'ramdisk_name', 'board_id'])
309
310    def __init__(self):
311        self.entries = []
312        self.ramdisk_total_size = 0
313        self.ramdisk_names = set()
314
315    def add_entry(self, ramdisk_path, ramdisk_type, ramdisk_name, board_id):
316        # Strip any trailing null for simple comparison.
317        stripped_ramdisk_name = ramdisk_name.rstrip(b'\x00')
318        if stripped_ramdisk_name in VENDOR_RAMDISK_NAME_BLOCKLIST:
319            raise ValueError(
320                f'Banned vendor ramdisk name: {stripped_ramdisk_name}')
321        if stripped_ramdisk_name in self.ramdisk_names:
322            raise ValueError(
323                f'Duplicated vendor ramdisk name: {stripped_ramdisk_name}')
324        self.ramdisk_names.add(stripped_ramdisk_name)
325
326        if board_id is None:
327            board_id = array.array(
328                'I', [0] * VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE)
329        else:
330            board_id = array.array('I', board_id)
331        if len(board_id) != VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE:
332            raise ValueError('board_id size must be '
333                             f'{VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE}')
334
335        with open(ramdisk_path, 'rb') as f:
336            ramdisk_size = filesize(f)
337        self.entries.append(self.VendorRamdiskTableEntry(
338            ramdisk_path, ramdisk_size, self.ramdisk_total_size, ramdisk_type,
339            ramdisk_name, board_id))
340        self.ramdisk_total_size += ramdisk_size
341
342    def write_ramdisks_padded(self, fout, alignment):
343        for entry in self.entries:
344            with open(entry.ramdisk_path, 'rb') as f:
345                fout.write(f.read())
346        pad_file(fout, alignment)
347
348    def write_entries_padded(self, fout, alignment):
349        for entry in self.entries:
350            fout.write(pack('I', entry.ramdisk_size))
351            fout.write(pack('I', entry.ramdisk_offset))
352            fout.write(pack('I', entry.ramdisk_type))
353            fout.write(pack(f'{VENDOR_RAMDISK_NAME_SIZE}s',
354                            entry.ramdisk_name))
355            fout.write(entry.board_id)
356        pad_file(fout, alignment)
357
358
359def write_padded_file(f_out, f_in, padding):
360    if f_in is None:
361        return
362    f_out.write(f_in.read())
363    pad_file(f_out, padding)
364
365
366def parse_int(x):
367    return int(x, 0)
368
369
370def parse_os_version(x):
371    match = re.search(r'^(\d{1,3})(?:\.(\d{1,3})(?:\.(\d{1,3}))?)?', x)
372    if match:
373        a = int(match.group(1))
374        b = c = 0
375        if match.lastindex >= 2:
376            b = int(match.group(2))
377        if match.lastindex == 3:
378            c = int(match.group(3))
379        # 7 bits allocated for each field
380        assert a < 128
381        assert b < 128
382        assert c < 128
383        return (a << 14) | (b << 7) | c
384    return 0
385
386
387def parse_os_patch_level(x):
388    match = re.search(r'^(\d{4})-(\d{2})(?:-(\d{2}))?', x)
389    if match:
390        y = int(match.group(1)) - 2000
391        m = int(match.group(2))
392        # 7 bits allocated for the year, 4 bits for the month
393        assert 0 <= y < 128
394        assert 0 < m <= 12
395        return (y << 4) | m
396    return 0
397
398
399def parse_vendor_ramdisk_type(x):
400    type_dict = {
401        'none': VENDOR_RAMDISK_TYPE_NONE,
402        'platform': VENDOR_RAMDISK_TYPE_PLATFORM,
403        'recovery': VENDOR_RAMDISK_TYPE_RECOVERY,
404        'dlkm': VENDOR_RAMDISK_TYPE_DLKM,
405    }
406    if x.lower() in type_dict:
407        return type_dict[x.lower()]
408    return parse_int(x)
409
410
411def get_vendor_boot_v4_usage():
412    return """vendor boot version 4 arguments:
413  --ramdisk_type {none,platform,recovery,dlkm}
414                        specify the type of the ramdisk
415  --ramdisk_name NAME
416                        specify the name of the ramdisk
417  --board_id{0..15} NUMBER
418                        specify the value of the board_id vector, defaults to 0
419  --vendor_ramdisk_fragment VENDOR_RAMDISK_FILE
420                        path to the vendor ramdisk file
421
422  These options can be specified multiple times, where each vendor ramdisk
423  option group ends with a --vendor_ramdisk_fragment option.
424  Each option group appends an additional ramdisk to the vendor boot image.
425"""
426
427
428def parse_vendor_ramdisk_args(args, args_list):
429    """Parses vendor ramdisk specific arguments.
430
431    Args:
432        args: An argparse.Namespace object. Parsed results are stored into this
433            object.
434        args_list: A list of argument strings to be parsed.
435
436    Returns:
437        A list argument strings that are not parsed by this method.
438    """
439    parser = ArgumentParser(add_help=False)
440    parser.add_argument('--ramdisk_type', type=parse_vendor_ramdisk_type,
441                        default=VENDOR_RAMDISK_TYPE_NONE)
442    parser.add_argument('--ramdisk_name',
443                        type=AsciizBytes(bufsize=VENDOR_RAMDISK_NAME_SIZE),
444                        required=True)
445    for i in range(VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE):
446        parser.add_argument(f'--board_id{i}', type=parse_int, default=0)
447    parser.add_argument(PARSER_ARGUMENT_VENDOR_RAMDISK_FRAGMENT, required=True)
448
449    unknown_args = []
450
451    vendor_ramdisk_table_builder = VendorRamdiskTableBuilder()
452    if args.vendor_ramdisk is not None:
453        vendor_ramdisk_table_builder.add_entry(
454            args.vendor_ramdisk.name, VENDOR_RAMDISK_TYPE_PLATFORM, b'', None)
455
456    while PARSER_ARGUMENT_VENDOR_RAMDISK_FRAGMENT in args_list:
457        idx = args_list.index(PARSER_ARGUMENT_VENDOR_RAMDISK_FRAGMENT) + 2
458        vendor_ramdisk_args = args_list[:idx]
459        args_list = args_list[idx:]
460
461        ramdisk_args, extra_args = parser.parse_known_args(vendor_ramdisk_args)
462        ramdisk_args_dict = vars(ramdisk_args)
463        unknown_args.extend(extra_args)
464
465        ramdisk_path = ramdisk_args.vendor_ramdisk_fragment
466        ramdisk_type = ramdisk_args.ramdisk_type
467        ramdisk_name = ramdisk_args.ramdisk_name
468        board_id = [ramdisk_args_dict[f'board_id{i}']
469                    for i in range(VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE)]
470        vendor_ramdisk_table_builder.add_entry(ramdisk_path, ramdisk_type,
471                                               ramdisk_name, board_id)
472
473    if len(args_list) > 0:
474        unknown_args.extend(args_list)
475
476    args.vendor_ramdisk_total_size = (vendor_ramdisk_table_builder
477                                      .ramdisk_total_size)
478    args.vendor_ramdisk_table_entry_num = len(vendor_ramdisk_table_builder
479                                              .entries)
480    args.vendor_ramdisk_table_builder = vendor_ramdisk_table_builder
481    return unknown_args
482
483
484def parse_cmdline():
485    version_parser = ArgumentParser(add_help=False)
486    version_parser.add_argument('--header_version', type=parse_int, default=0)
487    if version_parser.parse_known_args()[0].header_version < 3:
488        # For boot header v0 to v2, the kernel commandline field is split into
489        # two fields, cmdline and extra_cmdline. Both fields are asciiz strings,
490        # so we minus one here to ensure the encoded string plus the
491        # null-terminator can fit in the buffer size.
492        cmdline_size = BOOT_ARGS_SIZE + BOOT_EXTRA_ARGS_SIZE - 1
493    else:
494        cmdline_size = BOOT_ARGS_SIZE + BOOT_EXTRA_ARGS_SIZE
495
496    parser = ArgumentParser(formatter_class=RawDescriptionHelpFormatter,
497                            epilog=get_vendor_boot_v4_usage())
498    parser.add_argument('--kernel', type=FileType('rb'),
499                        help='path to the kernel')
500    parser.add_argument('--ramdisk', type=FileType('rb'),
501                        help='path to the ramdisk')
502    parser.add_argument('--second', type=FileType('rb'),
503                        help='path to the second bootloader')
504    parser.add_argument('--dtb', type=FileType('rb'), help='path to the dtb')
505    dtbo_group = parser.add_mutually_exclusive_group()
506    dtbo_group.add_argument('--recovery_dtbo', type=FileType('rb'),
507                            help='path to the recovery DTBO')
508    dtbo_group.add_argument('--recovery_acpio', type=FileType('rb'),
509                            metavar='RECOVERY_ACPIO', dest='recovery_dtbo',
510                            help='path to the recovery ACPIO')
511    parser.add_argument('--cmdline', type=AsciizBytes(bufsize=cmdline_size),
512                        default='', help='kernel command line arguments')
513    parser.add_argument('--vendor_cmdline',
514                        type=AsciizBytes(bufsize=VENDOR_BOOT_ARGS_SIZE),
515                        default='',
516                        help='vendor boot kernel command line arguments')
517    parser.add_argument('--base', type=parse_int, default=0x10000000,
518                        help='base address')
519    parser.add_argument('--kernel_offset', type=parse_int, default=0x00008000,
520                        help='kernel offset')
521    parser.add_argument('--ramdisk_offset', type=parse_int, default=0x01000000,
522                        help='ramdisk offset')
523    parser.add_argument('--second_offset', type=parse_int, default=0x00f00000,
524                        help='second bootloader offset')
525    parser.add_argument('--dtb_offset', type=parse_int, default=0x01f00000,
526                        help='dtb offset')
527
528    parser.add_argument('--os_version', type=parse_os_version, default=0,
529                        help='operating system version')
530    parser.add_argument('--os_patch_level', type=parse_os_patch_level,
531                        default=0, help='operating system patch level')
532    parser.add_argument('--tags_offset', type=parse_int, default=0x00000100,
533                        help='tags offset')
534    parser.add_argument('--board', type=AsciizBytes(bufsize=BOOT_NAME_SIZE),
535                        default='', help='board name')
536    parser.add_argument('--pagesize', type=parse_int,
537                        choices=[2**i for i in range(11, 15)], default=2048,
538                        help='page size')
539    parser.add_argument('--id', action='store_true',
540                        help='print the image ID on standard output')
541    parser.add_argument('--header_version', type=parse_int, default=0,
542                        help='boot image header version')
543    parser.add_argument('-o', '--output', type=FileType('wb'),
544                        help='output file name')
545    parser.add_argument('--vendor_boot', type=FileType('wb'),
546                        help='vendor boot output file name')
547    parser.add_argument('--vendor_ramdisk', type=FileType('rb'),
548                        help='path to the vendor ramdisk')
549    parser.add_argument('--vendor_bootconfig', type=FileType('rb'),
550                        help='path to the vendor bootconfig file')
551
552    gki_2_0_signing_args = parser.add_argument_group(
553        '[DEPRECATED] GKI 2.0 signing arguments')
554    gki_2_0_signing_args.add_argument(
555        '--gki_signing_algorithm', help='GKI signing algorithm to use')
556    gki_2_0_signing_args.add_argument(
557        '--gki_signing_key', help='path to RSA private key file')
558    gki_2_0_signing_args.add_argument(
559        '--gki_signing_signature_args', default='',
560        help='other hash arguments passed to avbtool')
561    gki_2_0_signing_args.add_argument(
562        '--gki_signing_avbtool_path', default='avbtool',
563        help='path to avbtool for boot signature generation')
564
565    args, extra_args = parser.parse_known_args()
566    if args.vendor_boot is not None and args.header_version > 3:
567        extra_args = parse_vendor_ramdisk_args(args, extra_args)
568    if len(extra_args) > 0:
569        raise ValueError(f'Unrecognized arguments: {extra_args}')
570
571    if args.header_version < 3:
572        args.extra_cmdline = args.cmdline[BOOT_ARGS_SIZE-1:]
573        args.cmdline = args.cmdline[:BOOT_ARGS_SIZE-1] + b'\x00'
574        assert len(args.cmdline) <= BOOT_ARGS_SIZE
575        assert len(args.extra_cmdline) <= BOOT_EXTRA_ARGS_SIZE
576
577    return args
578
579
580def add_boot_image_signature(args, pagesize):
581    """Adds the boot image signature.
582
583    Note that the signature will only be verified in VTS to ensure a
584    generic boot.img is used. It will not be used by the device
585    bootloader at boot time. The bootloader should only verify
586    the boot vbmeta at the end of the boot partition (or in the top-level
587    vbmeta partition) via the Android Verified Boot process, when the
588    device boots.
589    """
590    # Flush the buffer for signature calculation.
591    args.output.flush()
592
593    # Outputs the signed vbmeta to a separate file, then append to boot.img
594    # as the boot signature.
595    with tempfile.TemporaryDirectory() as temp_out_dir:
596        boot_signature_output = os.path.join(temp_out_dir, 'boot_signature')
597        generate_gki_certificate(
598            image=args.output.name, avbtool=args.gki_signing_avbtool_path,
599            name='boot', algorithm=args.gki_signing_algorithm,
600            key=args.gki_signing_key, salt='d00df00d',
601            additional_avb_args=args.gki_signing_signature_args.split(),
602            output=boot_signature_output,
603        )
604        with open(boot_signature_output, 'rb') as boot_signature:
605            boot_signature_bytes = boot_signature.read()
606            if len(boot_signature_bytes) > BOOT_IMAGE_V4_SIGNATURE_SIZE:
607                raise ValueError(
608                    f'boot sigature size is > {BOOT_IMAGE_V4_SIGNATURE_SIZE}')
609            boot_signature_bytes += b'\x00' * (
610                BOOT_IMAGE_V4_SIGNATURE_SIZE - len(boot_signature_bytes))
611            assert len(boot_signature_bytes) == BOOT_IMAGE_V4_SIGNATURE_SIZE
612            args.output.write(boot_signature_bytes)
613            pad_file(args.output, pagesize)
614
615
616def write_data(args, pagesize):
617    write_padded_file(args.output, args.kernel, pagesize)
618    write_padded_file(args.output, args.ramdisk, pagesize)
619    write_padded_file(args.output, args.second, pagesize)
620
621    if args.header_version > 0 and args.header_version < 3:
622        write_padded_file(args.output, args.recovery_dtbo, pagesize)
623    if args.header_version == 2:
624        write_padded_file(args.output, args.dtb, pagesize)
625    if args.header_version >= 4 and should_add_legacy_gki_boot_signature(args):
626        add_boot_image_signature(args, pagesize)
627
628
629def write_vendor_boot_data(args):
630    if args.header_version > 3:
631        builder = args.vendor_ramdisk_table_builder
632        builder.write_ramdisks_padded(args.vendor_boot, args.pagesize)
633        write_padded_file(args.vendor_boot, args.dtb, args.pagesize)
634        builder.write_entries_padded(args.vendor_boot, args.pagesize)
635        write_padded_file(args.vendor_boot, args.vendor_bootconfig,
636            args.pagesize)
637    else:
638        write_padded_file(args.vendor_boot, args.vendor_ramdisk, args.pagesize)
639        write_padded_file(args.vendor_boot, args.dtb, args.pagesize)
640
641
642def main():
643    args = parse_cmdline()
644    if args.vendor_boot is not None:
645        if args.header_version not in {3, 4}:
646            raise ValueError(
647                '--vendor_boot not compatible with given header version')
648        if args.header_version == 3 and args.vendor_ramdisk is None:
649            raise ValueError('--vendor_ramdisk missing or invalid')
650        write_vendor_boot_header(args)
651        write_vendor_boot_data(args)
652    if args.output is not None:
653        if args.second is not None and args.header_version > 2:
654            raise ValueError(
655                '--second not compatible with given header version')
656        img_id = write_header(args)
657        if args.header_version > 2:
658            write_data(args, BOOT_IMAGE_HEADER_V3_PAGESIZE)
659        else:
660            write_data(args, args.pagesize)
661        if args.id and img_id is not None:
662            print('0x' + ''.join(f'{octet:02x}' for octet in img_id))
663
664
665if __name__ == '__main__':
666    main()
667