1 files changed, 115 insertions, 0 deletions
diff --git a/tools/mcuboot/imgtool/keys/rsa_test.py b/tools/mcuboot/imgtool/keys/rsa_test.py
new file mode 100644
index 00000000..b0afa835
--- /dev/null
+++ b/tools/mcuboot/imgtool/keys/rsa_test.py
@@ -0,0 +1,115 @@
+"""
+Tests for RSA keys
+"""
+
+import io
+import os
+import sys
+import tempfile
+import unittest
+
+from cryptography.exceptions import InvalidSignature
+from cryptography.hazmat.primitives.asymmetric.padding import PSS, MGF1
+from cryptography.hazmat.primitives.hashes import SHA256
+
+# Setup sys path so 'imgtool' is in it.
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__),
+                                                '../..')))
+
+from imgtool.keys import load, RSA, RSAUsageError
+from imgtool.keys.rsa import RSA_KEY_SIZES
+
+
+class KeyGeneration(unittest.TestCase):
+
+    def setUp(self):
+        self.test_dir = tempfile.TemporaryDirectory()
+
+    def tname(self, base):
+        return os.path.join(self.test_dir.name, base)
+
+    def tearDown(self):
+        self.test_dir.cleanup()
+
+    def test_keygen(self):
+        # Try generating a RSA key with non-supported size
+        with self.assertRaises(RSAUsageError):
+            RSA.generate(key_size=1024)
+
+        for key_size in RSA_KEY_SIZES:
+            name1 = self.tname("keygen.pem")
+            k = RSA.generate(key_size=key_size)
+            k.export_private(name1, b'secret')
+
+            # Try loading the key without a password.
+            self.assertIsNone(load(name1))
+
+            k2 = load(name1, b'secret')
+
+            pubname = self.tname('keygen-pub.pem')
+            k2.export_public(pubname)
+            pk2 = load(pubname)
+
+            # We should be able to export the public key from the loaded
+            # public key, but not the private key.
+            pk2.export_public(self.tname('keygen-pub2.pem'))
+            self.assertRaises(RSAUsageError, pk2.export_private,
+                              self.tname('keygen-priv2.pem'))
+
+    def test_emit(self):
+        """Basic sanity check on the code emitters."""
+        for key_size in RSA_KEY_SIZES:
+            k = RSA.generate(key_size=key_size)
+
+            ccode = io.StringIO()
+            k.emit_c_public(ccode)
+            self.assertIn("rsa_pub_key", ccode.getvalue())
+            self.assertIn("rsa_pub_key_len", ccode.getvalue())
+
+            rustcode = io.StringIO()
+            k.emit_rust_public(rustcode)
+            self.assertIn("RSA_PUB_KEY", rustcode.getvalue())
+
+    def test_emit_pub(self):
+        """Basic sanity check on the code emitters, from public key."""
+        pubname = self.tname("public.pem")
+        for key_size in RSA_KEY_SIZES:
+            k = RSA.generate(key_size=key_size)
+            k.export_public(pubname)
+
+            k2 = load(pubname)
+
+            ccode = io.StringIO()
+            k2.emit_c_public(ccode)
+            self.assertIn("rsa_pub_key", ccode.getvalue())
+            self.assertIn("rsa_pub_key_len", ccode.getvalue())
+
+            rustcode = io.StringIO()
+            k2.emit_rust_public(rustcode)
+            self.assertIn("RSA_PUB_KEY", rustcode.getvalue())
+
+    def test_sig(self):
+        for key_size in RSA_KEY_SIZES:
+            k = RSA.generate(key_size=key_size)
+            buf = b'This is the message'
+            sig = k.sign(buf)
+
+            # The code doesn't have any verification, so verify this
+            # manually.
+            k.key.public_key().verify(
+                signature=sig,
+                data=buf,
+                padding=PSS(mgf=MGF1(SHA256()), salt_length=32),
+                algorithm=SHA256())
+
+            # Modify the message to make sure the signature fails.
+            self.assertRaises(InvalidSignature,
+                              k.key.public_key().verify,
+                              signature=sig,
+                              data=b'This is thE message',
+                              padding=PSS(mgf=MGF1(SHA256()), salt_length=32),
+                              algorithm=SHA256())
+
+
+if __name__ == '__main__':
+    unittest.main()