1"""
2Tests for ECDSA keys
3"""
4
5# SPDX-License-Identifier: Apache-2.0
6
7import io
8import os.path
9import sys
10import tempfile
11import unittest
12
13from cryptography.exceptions import InvalidSignature
14from cryptography.hazmat.primitives.asymmetric import ec
15from cryptography.hazmat.primitives.hashes import SHA256
16
17sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
18
19from imgtool.keys import load, ECDSA256P1, ECDSAUsageError
20
21class EcKeyGeneration(unittest.TestCase):
22
23    def setUp(self):
24        self.test_dir = tempfile.TemporaryDirectory()
25
26    def tname(self, base):
27        return os.path.join(self.test_dir.name, base)
28
29    def tearDown(self):
30        self.test_dir.cleanup()
31
32    def test_keygen(self):
33        name1 = self.tname("keygen.pem")
34        k = ECDSA256P1.generate()
35        k.export_private(name1, b'secret')
36
37        self.assertIsNone(load(name1))
38
39        k2 = load(name1, b'secret')
40
41        pubname = self.tname('keygen-pub.pem')
42        k2.export_public(pubname)
43        pk2 = load(pubname)
44
45        # We should be able to export the public key from the loaded
46        # public key, but not the private key.
47        pk2.export_public(self.tname('keygen-pub2.pem'))
48        self.assertRaises(ECDSAUsageError,
49                pk2.export_private, self.tname('keygen-priv2.pem'))
50
51    def test_emit(self):
52        """Basic sanity check on the code emitters."""
53        k = ECDSA256P1.generate()
54
55        ccode = io.StringIO()
56        k.emit_c_public(ccode)
57        self.assertIn("ecdsa_pub_key", ccode.getvalue())
58        self.assertIn("ecdsa_pub_key_len", ccode.getvalue())
59
60        rustcode = io.StringIO()
61        k.emit_rust_public(rustcode)
62        self.assertIn("ECDSA_PUB_KEY", rustcode.getvalue())
63
64    def test_emit_pub(self):
65        """Basic sanity check on the code emitters."""
66        pubname = self.tname("public.pem")
67        k = ECDSA256P1.generate()
68        k.export_public(pubname)
69
70        k2 = load(pubname)
71
72        ccode = io.StringIO()
73        k2.emit_c_public(ccode)
74        self.assertIn("ecdsa_pub_key", ccode.getvalue())
75        self.assertIn("ecdsa_pub_key_len", ccode.getvalue())
76
77        rustcode = io.StringIO()
78        k2.emit_rust_public(rustcode)
79        self.assertIn("ECDSA_PUB_KEY", rustcode.getvalue())
80
81    def test_sig(self):
82        k = ECDSA256P1.generate()
83        buf = b'This is the message'
84        sig = k.raw_sign(buf)
85
86        # The code doesn't have any verification, so verify this
87        # manually.
88        k.key.public_key().verify(
89                signature=sig,
90                data=buf,
91                signature_algorithm=ec.ECDSA(SHA256()))
92
93        # Modify the message to make sure the signature fails.
94        self.assertRaises(InvalidSignature,
95                k.key.public_key().verify,
96                signature=sig,
97                data=b'This is thE message',
98                signature_algorithm=ec.ECDSA(SHA256()))
99
100if __name__ == '__main__':
101    unittest.main()
102