Add key generation unit test

dkim/tests/test_dkim_generate.py

@@ -0,0 +1,115 @@
+# This software is provided 'as-is', without any express or implied
+# warranty.  In no event will the author be held liable for any damages
+# arising from the use of this software.
+#
+# Permission is granted to anyone to use this software for any purpose,
+# including commercial applications, and to alter it and redistribute it
+# freely, subject to the following restrictions:
+#
+# 1. The origin of this software must not be misrepresented; you must not
+#    claim that you wrote the original software. If you use this software
+#    in a product, an acknowledgment in the product documentation would be
+#    appreciated but is not required.
+# 2. Altered source versions must be plainly marked as such, and must not be
+#    misrepresented as being the original software.
+# 3. This notice may not be removed or altered from any source distribution.
+#
+# Copyright (c) 2011 William Grant <me@williamgrant.id.au>
+# Copyright (c) 2022 Adrien Precigout <dev@asdrip.fr>
+
+import os.path
+import tempfile
+import unittest
+
+import dkim
+import dknewkey
+
+def read_test_data(filename):
+    """Get the content of the given test data file.
+
+    The files live in dkim/tests/data.
+    """
+    path = os.path.join(os.path.dirname(__file__), 'data', filename)
+    with open(path, 'rb') as f:
+        return f.read()
+
+
+class TestSignAndVerify(unittest.TestCase):
+    """End-to-end signature and verification tests with a generated key."""
+
+    def setUp(self):
+        self.message = read_test_data("test.message")
+        self.ed25519_dns_key_file = ""
+        self.rsa_dns_key_file = ""
+    
+
+    def test_generate_verifies_new_RSA_key(self):
+        #Create temporary dir
+        tmpdir = tempfile.TemporaryDirectory()
+        keydir = tmpdir.name
+        rsa_key_file = os.path.join(keydir, "dkim.rsa.key")
+        self.rsa_dns_key_file = os.path.join(keydir, "dkim.rsa.key.pub.txt")
+        #Generate a rsa key
+        dknewkey.GenRSAKeys(rsa_key_file, False)
+        dknewkey.ExtractRSADnsPublicKey(rsa_key_file, self.rsa_dns_key_file, False)
+        #Load the key
+        rsakey = read_test_data(rsa_key_file)
+        #Test signature with the newely generated key  
+        for header_algo in (b"simple", b"relaxed"):
+            for body_algo in (b"simple", b"relaxed"):
+                sig = dkim.sign(
+                    self.message, b"test", b"example.com", rsakey,
+                    canonicalize=(header_algo, body_algo))
+                res = dkim.verify(sig + self.message, dnsfunc=self.dnsfuncRSA)
+                self.assertTrue(res)
+
+
+    def test_generate_verifies_Ed25519_key(self):
+        #Create temporary dir
+        tmpdir = tempfile.TemporaryDirectory()
+        keydir = tmpdir.name
+        ed25519_key_file = os.path.join(keydir, "dkim.ed25519.key")
+        self.ed25519_dns_key_file = os.path.join(keydir, "dkim.ed25519.key.pub.txt")
+        #Generate a ed25519 key
+        pkt = dknewkey.GenEd25519Keys(ed25519_key_file, False)
+        dknewkey.ExtractEd25519PublicKey(self.ed25519_dns_key_file, pkt, False)
+        #Load the key
+        ed25519key = read_test_data(ed25519_key_file)
+        #Test signature with the newely generated key 
+        for header_algo in (b"simple", b"relaxed"):
+            for body_algo in (b"simple", b"relaxed"):
+                sig = dkim.sign(
+                    self.message, b"test1", b"example.com", ed25519key,
+                    signature_algorithm=b'ed25519-sha256',
+                    canonicalize=(header_algo, body_algo))
+                res = dkim.verify(sig + self.message, dnsfunc=self.dnsfuncED25519)
+                self.assertTrue(res)
+
+
+    def dnsfuncRSA(self, domain, timeout=5):
+        _dns_responses = {
+            'test._domainkey.example.com.': read_test_data(self.rsa_dns_key_file),
+        }
+        try:
+            domain = domain.decode('ascii')
+        except UnicodeDecodeError:
+            return None
+        self.assertTrue(domain in _dns_responses,domain)
+        return _dns_responses[domain]
+
+    def dnsfuncED25519(self, domain, timeout=5):
+        _dns_responses = {
+            'test1._domainkey.example.com.': read_test_data(self.ed25519_dns_key_file),
+        }
+        try:
+            domain = domain.decode('ascii')
+        except UnicodeDecodeError:
+            return None
+        self.assertTrue(domain in _dns_responses,domain)
+        return _dns_responses[domain]
+
+   
+
+def test_suite():
+    from unittest import TestLoader
+    return TestLoader().loadTestsFromName(__name__)