1# SPDX-License-Identifier:	GPL-2.0+
2#
3# Copyright (c) 2020,2021 Alexandru Gagniuc <mr.nuke.me@gmail.com>
4
5"""
6Test ECDSA signing of FIT images
7
8This test uses mkimage to sign an existing FIT image with an ECDSA key. The
9signature is then extracted, and verified against pyCryptodome.
10This test doesn't run the sandbox. It only checks the host tool 'mkimage'
11"""
12
13import os
14import pytest
15import utils
16from Cryptodome.Hash import SHA256
17from Cryptodome.PublicKey import ECC
18from Cryptodome.Signature import DSS
19
20class SignableFitImage(object):
21    """ Helper to manipulate a FIT image on disk """
22    def __init__(self, ubman, file_name):
23        self.fit = file_name
24        self.ubman = ubman
25        self.signable_nodes = set()
26
27    def __fdt_list(self, path):
28        return utils.run_and_log(self.ubman, f'fdtget -l {self.fit} {path}')
29
30    def __fdt_set(self, node, **prop_value):
31        for prop, value in prop_value.items():
32            utils.run_and_log(self.ubman,
33                              f'fdtput -ts {self.fit} {node} {prop} {value}')
34
35    def __fdt_get_binary(self, node, prop):
36        numbers = utils.run_and_log(self.ubman,
37                                    f'fdtget -tbi {self.fit} {node} {prop}')
38
39        bignum = bytearray()
40        for little_num in numbers.split():
41            bignum.append(int(little_num))
42
43        return bignum
44
45    def find_signable_image_nodes(self):
46        for node in self.__fdt_list('/images').split():
47            image = f'/images/{node}'
48            if 'signature' in self.__fdt_list(image):
49                self.signable_nodes.add(image)
50
51        return self.signable_nodes
52
53    def change_signature_algo_to_ecdsa(self):
54        for image in self.signable_nodes:
55            self.__fdt_set(f'{image}/signature', algo='sha256,ecdsa256')
56
57    def sign(self, mkimage, key_file):
58        utils.run_and_log(self.ubman, [mkimage, '-F', self.fit, f'-G{key_file}'])
59
60    def check_signatures(self, key):
61        for image in self.signable_nodes:
62            raw_sig = self.__fdt_get_binary(f'{image}/signature', 'value')
63            raw_bin = self.__fdt_get_binary(image, 'data')
64
65            sha = SHA256.new(raw_bin)
66            verifier = DSS.new(key, 'fips-186-3')
67            verifier.verify(sha, bytes(raw_sig))
68
69
70@pytest.mark.buildconfigspec('fit_signature')
71@pytest.mark.requiredtool('dtc')
72@pytest.mark.requiredtool('fdtget')
73@pytest.mark.requiredtool('fdtput')
74def test_fit_ecdsa(ubman):
75    """ Test that signatures generated by mkimage are legible. """
76    def generate_ecdsa_key():
77        return ECC.generate(curve='prime256v1')
78
79    def assemble_fit_image(dest_fit, its, destdir):
80        dtc_args = f'-I dts -O dtb -i {destdir}'
81        utils.run_and_log(ubman, [mkimage, '-D', dtc_args, '-f', its, dest_fit])
82
83    def dtc(dts):
84        dtb = dts.replace('.dts', '.dtb')
85        utils.run_and_log(ubman, f'dtc {datadir}/{dts} -O dtb -o {tempdir}/{dtb}')
86
87    mkimage = ubman.config.build_dir + '/tools/mkimage'
88    datadir = ubman.config.source_dir + '/test/py/tests/vboot/'
89    tempdir = os.path.join(ubman.config.result_dir, 'ecdsa')
90    os.makedirs(tempdir, exist_ok=True)
91    key_file = f'{tempdir}/ecdsa-test-key.pem'
92    fit_file = f'{tempdir}/test.fit'
93    dtc('sandbox-kernel.dts')
94
95    key = generate_ecdsa_key()
96
97    # Create a fake kernel image -- zeroes will do just fine
98    with open(f'{tempdir}/test-kernel.bin', 'w') as fd:
99        fd.write(500 * chr(0))
100
101    # invocations of mkimage expect to read the key from disk
102    with open(key_file, 'w') as f:
103        f.write(key.export_key(format='PEM'))
104
105    assemble_fit_image(fit_file, f'{datadir}/sign-images-sha256.its', tempdir)
106
107    fit = SignableFitImage(ubman, fit_file)
108    nodes = fit.find_signable_image_nodes()
109    if len(nodes) == 0:
110        raise ValueError('FIT image has no "/image" nodes with "signature"')
111
112    fit.change_signature_algo_to_ecdsa()
113    fit.sign(mkimage, key_file)
114    fit.check_signatures(key)
115