1#!/usr/bin/env python3
2#
3# Copyright 2022, 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
18"""Certify a GKI boot image by generating and appending its boot_signature."""
19
20from argparse import ArgumentParser
21import glob
22import os
23import shlex
24import shutil
25import subprocess
26import tempfile
27
28from gki.generate_gki_certificate import generate_gki_certificate
29from unpack_bootimg import unpack_bootimg
30
31BOOT_SIGNATURE_SIZE = 16 * 1024
32
33
34def get_kernel(boot_img):
35    """Extracts the kernel from |boot_img| and returns it."""
36    with tempfile.TemporaryDirectory() as unpack_dir:
37        unpack_bootimg(boot_img, unpack_dir)
38        with open(os.path.join(unpack_dir, 'kernel'), 'rb') as kernel:
39            kernel_bytes = kernel.read()
40            assert len(kernel_bytes) > 0
41            return kernel_bytes
42
43
44def add_certificate(boot_img, algorithm, key, extra_args):
45    """Appends certificates to the end of the boot image.
46
47    This functions appends two certificates to the end of the |boot_img|:
48    the 'boot' certificate and the 'generic_kernel' certificate. The former
49    is to certify the entire |boot_img|, while the latter is to certify
50    the kernel inside the |boot_img|.
51    """
52
53    def generate_certificate(image, certificate_name):
54        """Generates the certificate and returns the certificate content."""
55        with tempfile.NamedTemporaryFile() as output_certificate:
56            generate_gki_certificate(
57                image=image, avbtool='avbtool', name=certificate_name,
58                algorithm=algorithm, key=key, salt='d00df00d',
59                additional_avb_args=extra_args, output=output_certificate.name)
60            output_certificate.seek(os.SEEK_SET, 0)
61            return output_certificate.read()
62
63    boot_signature_bytes = b''
64    boot_signature_bytes += generate_certificate(boot_img, 'boot')
65
66    with tempfile.NamedTemporaryFile() as kernel_img:
67        kernel_img.write(get_kernel(boot_img))
68        kernel_img.flush()
69        boot_signature_bytes += generate_certificate(kernel_img.name,
70                                                     'generic_kernel')
71
72    if len(boot_signature_bytes) > BOOT_SIGNATURE_SIZE:
73        raise ValueError(
74            f'boot_signature size must be <= {BOOT_SIGNATURE_SIZE}')
75    boot_signature_bytes += (
76        b'\0' * (BOOT_SIGNATURE_SIZE - len(boot_signature_bytes)))
77    assert len(boot_signature_bytes) == BOOT_SIGNATURE_SIZE
78
79    with open(boot_img, 'ab') as f:
80        f.write(boot_signature_bytes)
81
82
83def erase_certificate_and_avb_footer(boot_img):
84    """Erases the boot certificate and avb footer.
85
86    A boot image might already contain a certificate and/or a AVB footer.
87    This function erases these additional metadata from the |boot_img|.
88    """
89    # Tries to erase the AVB footer first, which may or may not exist.
90    avbtool_cmd = ['avbtool', 'erase_footer', '--image', boot_img]
91    subprocess.run(avbtool_cmd, check=False, stderr=subprocess.DEVNULL)
92    assert os.path.getsize(boot_img) > 0
93
94    # No boot signature to erase, just return.
95    if os.path.getsize(boot_img) <= BOOT_SIGNATURE_SIZE:
96        return
97
98    # Checks if the last 16K is a boot signature, then erases it.
99    with open(boot_img, 'rb') as image:
100        image.seek(-BOOT_SIGNATURE_SIZE, os.SEEK_END)
101        boot_signature = image.read(BOOT_SIGNATURE_SIZE)
102        assert len(boot_signature) == BOOT_SIGNATURE_SIZE
103
104    with tempfile.NamedTemporaryFile() as signature_tmpfile:
105        signature_tmpfile.write(boot_signature)
106        signature_tmpfile.flush()
107        avbtool_info_cmd = [
108            'avbtool', 'info_image', '--image', signature_tmpfile.name]
109        result = subprocess.run(avbtool_info_cmd, check=False,
110                                stdout=subprocess.DEVNULL,
111                                stderr=subprocess.DEVNULL)
112        has_boot_signature = (result.returncode == 0)
113
114    if has_boot_signature:
115        new_file_size = os.path.getsize(boot_img) - BOOT_SIGNATURE_SIZE
116        os.truncate(boot_img, new_file_size)
117
118    assert os.path.getsize(boot_img) > 0
119
120
121def get_avb_image_size(image):
122    """Returns the image size if there is a AVB footer, else return zero."""
123
124    avbtool_info_cmd = ['avbtool', 'info_image', '--image', image]
125    result = subprocess.run(avbtool_info_cmd, check=False,
126                            stdout=subprocess.DEVNULL,
127                            stderr=subprocess.DEVNULL)
128
129    if result.returncode == 0:
130        return os.path.getsize(image)
131
132    return 0
133
134
135def add_avb_footer(image, partition_size, extra_footer_args):
136    """Appends a AVB hash footer to the image."""
137
138    avbtool_cmd = ['avbtool', 'add_hash_footer', '--image', image,
139                   '--partition_name', 'boot']
140
141    if partition_size:
142        avbtool_cmd.extend(['--partition_size', str(partition_size)])
143    else:
144        avbtool_cmd.extend(['--dynamic_partition_size'])
145
146    avbtool_cmd.extend(extra_footer_args)
147    subprocess.check_call(avbtool_cmd)
148
149
150def load_dict_from_file(path):
151    """Loads key=value pairs from |path| and returns a dict."""
152    d = {}
153    with open(path, 'r', encoding='utf-8') as f:
154        for line in f:
155            line = line.strip()
156            if not line or line.startswith('#'):
157                continue
158            if '=' in line:
159                name, value = line.split('=', 1)
160                d[name] = value
161    return d
162
163
164def load_gki_info_file(gki_info_file, extra_args, extra_footer_args):
165    """Loads extra arguments from the gki info file.
166
167    Args:
168        gki_info_file: path to a gki-info.txt.
169        extra_args: the extra arguments forwarded to avbtool when creating
170          the gki certificate.
171        extra_footer_args: the extra arguments forwarded to avbtool when
172          creating the avb footer.
173
174    """
175    info_dict = load_dict_from_file(gki_info_file)
176    if 'certify_bootimg_extra_args' in info_dict:
177        extra_args.extend(
178            shlex.split(info_dict['certify_bootimg_extra_args']))
179    if 'certify_bootimg_extra_footer_args' in info_dict:
180        extra_footer_args.extend(
181            shlex.split(info_dict['certify_bootimg_extra_footer_args']))
182
183
184def get_archive_name_and_format_for_shutil(path):
185    """Returns archive name and format to shutil.make_archive() for the |path|.
186
187    e.g., returns ('/path/to/boot-img', 'gztar') if |path| is
188    '/path/to/boot-img.tar.gz'.
189    """
190    for format_name, format_extensions, _ in shutil.get_unpack_formats():
191        for extension in format_extensions:
192            if path.endswith(extension):
193                return path[:-len(extension)], format_name
194
195    raise ValueError(f"Unsupported archive format: '{path}'")
196
197
198def parse_cmdline():
199    """Parse command-line options."""
200    parser = ArgumentParser(add_help=True)
201
202    # Required args.
203    input_group = parser.add_mutually_exclusive_group(required=True)
204    input_group.add_argument(
205        '--boot_img', help='path to the boot image to certify')
206    input_group.add_argument(
207        '--boot_img_archive', help='path to the boot images archive to certify')
208
209    parser.add_argument('--algorithm', required=True,
210                        help='signing algorithm for the certificate')
211    parser.add_argument('--key', required=True,
212                        help='path to the RSA private key')
213    parser.add_argument('--gki_info',
214                        help='path to a gki-info.txt to append additional'
215                             'properties into the boot signature')
216    parser.add_argument('-o', '--output', required=True,
217                        help='output file name')
218
219    # Optional args.
220    parser.add_argument('--extra_args', default=[], action='append',
221                        help='extra arguments to be forwarded to avbtool')
222    parser.add_argument('--extra_footer_args', default=[], action='append',
223                        help='extra arguments for adding the avb footer')
224
225    args = parser.parse_args()
226
227    if args.gki_info and args.boot_img_archive:
228        parser.error('--gki_info cannot be used with --boot_image_archive. '
229                     'The gki_info file should be included in the archive.')
230
231    extra_args = []
232    for a in args.extra_args:
233        extra_args.extend(shlex.split(a))
234    args.extra_args = extra_args
235
236    extra_footer_args = []
237    for a in args.extra_footer_args:
238        extra_footer_args.extend(shlex.split(a))
239    args.extra_footer_args = extra_footer_args
240
241    if args.gki_info:
242        load_gki_info_file(args.gki_info,
243                           args.extra_args,
244                           args.extra_footer_args)
245
246    return args
247
248
249def certify_bootimg(boot_img, output_img, algorithm, key, extra_args,
250                    extra_footer_args):
251    """Certify a GKI boot image by generating and appending a boot_signature."""
252    with tempfile.TemporaryDirectory() as temp_dir:
253        boot_tmp = os.path.join(temp_dir, 'boot.tmp')
254        shutil.copy2(boot_img, boot_tmp)
255
256        erase_certificate_and_avb_footer(boot_tmp)
257        add_certificate(boot_tmp, algorithm, key, extra_args)
258
259        avb_partition_size = get_avb_image_size(boot_img)
260        add_avb_footer(boot_tmp, avb_partition_size, extra_footer_args)
261
262        # We're done, copy the temp image to the final output.
263        shutil.copy2(boot_tmp, output_img)
264
265
266def certify_bootimg_archive(boot_img_archive, output_archive,
267                            algorithm, key, extra_args, extra_footer_args):
268    """Similar to certify_bootimg(), but for an archive of boot images."""
269    with tempfile.TemporaryDirectory() as unpack_dir:
270        shutil.unpack_archive(boot_img_archive, unpack_dir)
271
272        gki_info_file = os.path.join(unpack_dir, 'gki-info.txt')
273        if os.path.exists(gki_info_file):
274            load_gki_info_file(gki_info_file, extra_args, extra_footer_args)
275
276        for boot_img in glob.glob(os.path.join(unpack_dir, 'boot*.img')):
277            print(f'Certifying {os.path.basename(boot_img)} ...')
278            certify_bootimg(boot_img=boot_img, output_img=boot_img,
279                            algorithm=algorithm, key=key, extra_args=extra_args,
280                            extra_footer_args=extra_footer_args)
281
282        print(f'Making certified archive: {output_archive}')
283        archive_file_name, archive_format = (
284            get_archive_name_and_format_for_shutil(output_archive))
285        built_archive = shutil.make_archive(archive_file_name,
286                                            archive_format,
287                                            unpack_dir)
288        # shutil.make_archive() builds *.tar.gz when then |archive_format| is
289        # 'gztar'. However, the end user might specify |output_archive| with
290        # *.tgz. Renaming *.tar.gz to *.tgz for this case.
291        if built_archive != os.path.realpath(output_archive):
292            print(f'Renaming {built_archive} -> {output_archive} ...')
293            os.rename(built_archive, output_archive)
294
295
296def main():
297    """Parse arguments and certify the boot image."""
298    args = parse_cmdline()
299
300    if args.boot_img_archive:
301        certify_bootimg_archive(args.boot_img_archive, args.output,
302                                args.algorithm, args.key, args.extra_args,
303                                args.extra_footer_args)
304    else:
305        certify_bootimg(args.boot_img, args.output, args.algorithm,
306                        args.key, args.extra_args, args.extra_footer_args)
307
308
309if __name__ == '__main__':
310    main()
311