[PATCH v4 6/6] test/py: ecdsa: Add test for mkimage ECDSA signing

[Alexandru Gagniuc <mr.nuke.me@gmail.com>]

Add a test to make sure that the ECDSA signatures generated by
mkimage can be verified successfully. pyCryptodomex was chosen as the
crypto library because it integrates much better with python code.
Using openssl would have been unnecessarily painful.

Signed-off-by: Alexandru Gagniuc <mr.nuke.me@gmail.com>
---
 test/py/tests/test_fit_ecdsa.py | 111 ++++++++++++++++++++++++++++++++
 1 file changed, 111 insertions(+)
 create mode 100644 test/py/tests/test_fit_ecdsa.py

diff --git a/test/py/tests/test_fit_ecdsa.py b/test/py/tests/test_fit_ecdsa.py
new file mode 100644
index 0000000000..6cadb0cbb5
--- /dev/null
+++ b/test/py/tests/test_fit_ecdsa.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier:	GPL-2.0+
+#
+# Copyright (c) 2020,2021 Alexandru Gagniuc <mr.nuke.me@gmail.com>
+
+"""
+Test ECDSA signing of FIT images
+
+This test uses mkimage to sign an existing FIT image with an ECDSA key. The
+signature is then extracted, and verified against pyCryptodome.
+This test doesn't run the sandbox. It only checks the host tool 'mkimage'
+"""
+
+import pytest
+import u_boot_utils as util
+from Cryptodome.Hash import SHA256
+from Cryptodome.PublicKey import ECC
+from Cryptodome.Signature import DSS
+
+class SignableFitImage(object):
+    """ Helper to manipulate a FIT image on disk """
+    def __init__(self, cons, file_name):
+        self.fit = file_name
+        self.cons = cons
+        self.signable_nodes = set()
+
+    def __fdt_list(self, path):
+        return util.run_and_log(self.cons, f'fdtget -l {self.fit} {path}')
+
+    def __fdt_set(self, node, **prop_value):
+        for prop, value in prop_value.items():
+            util.run_and_log(self.cons, f'fdtput -ts {self.fit} {node} {prop} {value}')
+
+    def __fdt_get_binary(self, node, prop):
+        numbers = util.run_and_log(self.cons, f'fdtget -tbi {self.fit} {node} {prop}')
+
+        bignum = bytearray()
+        for little_num in numbers.split():
+            bignum.append(int(little_num))
+
+        return bignum
+
+    def find_signable_image_nodes(self):
+        for node in self.__fdt_list('/images').split():
+            image = f'/images/{node}'
+            if 'signature' in self.__fdt_list(image):
+                self.signable_nodes.add(image)
+
+        return self.signable_nodes
+
+    def change_signature_algo_to_ecdsa(self):
+        for image in self.signable_nodes:
+            self.__fdt_set(f'{image}/signature', algo='sha256,ecdsa256')
+
+    def sign(self, mkimage, key_file):
+        util.run_and_log(self.cons, [mkimage, '-F', self.fit, f'-k{key_file}'])
+
+    def check_signatures(self, key):
+        for image in self.signable_nodes:
+            raw_sig = self.__fdt_get_binary(f'{image}/signature', 'value')
+            raw_bin = self.__fdt_get_binary(image, 'data')
+
+            sha = SHA256.new(raw_bin)
+            verifier = DSS.new(key, 'fips-186-3')
+            verifier.verify(sha, bytes(raw_sig))
+
+
+ at pytest.mark.buildconfigspec('fit_signature')
+ at pytest.mark.requiredtool('dtc')
+ at pytest.mark.requiredtool('fdtget')
+ at pytest.mark.requiredtool('fdtput')
+def test_fit_ecdsa(u_boot_console):
+    """ Test that signatures generated by mkimage are legible. """
+    def generate_ecdsa_key():
+        return ECC.generate(curve='prime256v1')
+
+    def assemble_fit_image(dest_fit, its, destdir):
+        dtc_args = f'-I dts -O dtb -i {destdir}'
+        util.run_and_log(cons, [mkimage, '-D', dtc_args, '-f', its, dest_fit])
+
+    def dtc(dts):
+        dtb = dts.replace('.dts', '.dtb')
+        util.run_and_log(cons, f'dtc {datadir}/{dts} -O dtb -o {tempdir}/{dtb}')
+
+    cons = u_boot_console
+    mkimage = cons.config.build_dir + '/tools/mkimage'
+    datadir = cons.config.source_dir + '/test/py/tests/vboot/'
+    tempdir = cons.config.result_dir
+    key_file = f'{tempdir}/ecdsa-test-key.pem'
+    fit_file = f'{tempdir}/test.fit'
+    dtc('sandbox-kernel.dts')
+
+    key = generate_ecdsa_key()
+
+    # Create a fake kernel image -- zeroes will do just fine
+    with open(f'{tempdir}/test-kernel.bin', 'w') as fd:
+        fd.write(500 * chr(0))
+
+    # invokations of mkimage expect to read the key from disk
+    with open(key_file, 'w') as f:
+        f.write(key.export_key(format='PEM'))
+
+    assemble_fit_image(fit_file, f'{datadir}/sign-images-sha256.its', tempdir)
+
+    fit = SignableFitImage(cons, fit_file)
+    nodes = fit.find_signable_image_nodes()
+    if len(nodes) == 0:
+        raise ValueError('FIT image has no "/image" nodes with "signature"')
+
+    fit.change_signature_algo_to_ecdsa()
+    fit.sign(mkimage, key_file)
+    fit.check_signatures(key)
-- 
2.26.2