1"""
2Tests for ECDSA keys
3"""
4
5# SPDX-License-Identifier: Apache-2.0
6
7import hashlib
8import io
9import os.path
10import sys
11import tempfile
12import unittest
13
14from cryptography.exceptions import InvalidSignature
15from cryptography.hazmat.primitives.asymmetric import ed25519
16
17sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
18
19from imgtool.keys import load, Ed25519, Ed25519UsageError
20
21
22class Ed25519KeyGeneration(unittest.TestCase):
23
24    def setUp(self):
25        self.test_dir = tempfile.TemporaryDirectory()
26
27    def tname(self, base):
28        return os.path.join(self.test_dir.name, base)
29
30    def tearDown(self):
31        self.test_dir.cleanup()
32
33    def test_keygen(self):
34        name1 = self.tname("keygen.pem")
35        k = Ed25519.generate()
36        k.export_private(name1, b'secret')
37
38        self.assertIsNone(load(name1))
39
40        k2 = load(name1, b'secret')
41
42        pubname = self.tname('keygen-pub.pem')
43        k2.export_public(pubname)
44        pk2 = load(pubname)
45
46        # We should be able to export the public key from the loaded
47        # public key, but not the private key.
48        pk2.export_public(self.tname('keygen-pub2.pem'))
49        self.assertRaises(Ed25519UsageError,
50                          pk2.export_private, self.tname('keygen-priv2.pem'))
51
52    def test_emit(self):
53        """Basic sanity check on the code emitters."""
54        k = Ed25519.generate()
55
56        pubpem = io.StringIO()
57        k.emit_public_pem(pubpem)
58        self.assertIn("BEGIN PUBLIC KEY", pubpem.getvalue())
59        self.assertIn("END PUBLIC KEY", pubpem.getvalue())
60
61        ccode = io.StringIO()
62        k.emit_c_public(ccode)
63        self.assertIn("ed25519_pub_key", ccode.getvalue())
64        self.assertIn("ed25519_pub_key_len", ccode.getvalue())
65
66        hashccode = io.StringIO()
67        k.emit_c_public_hash(hashccode)
68        self.assertIn("ed25519_pub_key_hash", hashccode.getvalue())
69        self.assertIn("ed25519_pub_key_hash_len", hashccode.getvalue())
70
71        rustcode = io.StringIO()
72        k.emit_rust_public(rustcode)
73        self.assertIn("ED25519_PUB_KEY", rustcode.getvalue())
74
75        # raw data - bytes
76        pubraw = io.BytesIO()
77        k.emit_raw_public(pubraw)
78        self.assertTrue(len(pubraw.getvalue()) > 0)
79
80        hashraw = io.BytesIO()
81        k.emit_raw_public_hash(hashraw)
82        self.assertTrue(len(hashraw.getvalue()) > 0)
83
84    def test_emit_pub(self):
85        """Basic sanity check on the code emitters, from public key."""
86        pubname = self.tname("public.pem")
87        k = Ed25519.generate()
88        k.export_public(pubname)
89
90        k2 = load(pubname)
91
92        ccode = io.StringIO()
93        k2.emit_c_public(ccode)
94        self.assertIn("ed25519_pub_key", ccode.getvalue())
95        self.assertIn("ed25519_pub_key_len", ccode.getvalue())
96
97        rustcode = io.StringIO()
98        k2.emit_rust_public(rustcode)
99        self.assertIn("ED25519_PUB_KEY", rustcode.getvalue())
100
101    def test_sig(self):
102        k = Ed25519.generate()
103        buf = b'This is the message'
104        sha = hashlib.sha256()
105        sha.update(buf)
106        digest = sha.digest()
107        sig = k.sign_digest(digest)
108
109        # The code doesn't have any verification, so verify this
110        # manually.
111        k.key.public_key().verify(signature=sig, data=digest)
112
113        # Modify the message to make sure the signature fails.
114        sha = hashlib.sha256()
115        sha.update(b'This is thE message')
116        new_digest = sha.digest()
117        self.assertRaises(InvalidSignature,
118                          k.key.public_key().verify,
119                          signature=sig,
120                          data=new_digest)
121
122
123if __name__ == '__main__':
124    unittest.main()
125