Improve consistency in DNSSEC code.

The make_ds method took its algorithm as a string, and the nsec3_hash
method took an algorithm as an int.  Change both of them to accept
either, and add enums for both sets of algorithms.

2 files changed, 37 insertions, 6 deletions

diff --git a/tests/test_dnssec.py b/tests/test_dnssec.py
index f38637d..c10becd 100644
--- a/tests/test_dnssec.py
+++ b/tests/test_dnssec.py
@@ -281,21 +281,28 @@ class DNSSECValidatorTestCase(unittest.TestCase):
 class DNSSECMakeDSTestCase(unittest.TestCase):
 
     def testMakeExampleSHA1DS(self):  # type: () -> None
-        ds = dns.dnssec.make_ds(abs_example, example_sep_key, 'SHA1')
-        self.assertEqual(ds, example_ds_sha1)
+        for algorithm in ('SHA1', 'sha1', dns.dnssec.DSDigest.SHA1):
+            ds = dns.dnssec.make_ds(abs_example, example_sep_key, algorithm)
+            self.assertEqual(ds, example_ds_sha1)
 
     def testMakeExampleSHA256DS(self):  # type: () -> None
-        ds = dns.dnssec.make_ds(abs_example, example_sep_key, 'SHA256')
-        self.assertEqual(ds, example_ds_sha256)
+        for algorithm in ('SHA256', 'sha256', dns.dnssec.DSDigest.SHA256):
+            ds = dns.dnssec.make_ds(abs_example, example_sep_key, algorithm)
+            self.assertEqual(ds, example_ds_sha256)
 
     def testMakeExampleSHA384DS(self):  # type: () -> None
-        ds = dns.dnssec.make_ds(abs_example, example_sep_key, 'SHA384')
-        self.assertEqual(ds, example_ds_sha384)
+        for algorithm in ('SHA384', 'sha384', dns.dnssec.DSDigest.SHA384):
+            ds = dns.dnssec.make_ds(abs_example, example_sep_key, algorithm)
+            self.assertEqual(ds, example_ds_sha384)
 
     def testMakeSHA256DS(self):  # type: () -> None
         ds = dns.dnssec.make_ds(abs_dnspython_org, sep_key, 'SHA256')
         self.assertEqual(ds, good_ds)
 
+    def testInvalidAlgorithm(self):  # type: () -> None
+        for algorithm in (10, 'shax'):
+            with self.assertRaises(dns.dnssec.UnsupportedAlgorithm):
+                ds = dns.dnssec.make_ds(abs_example, example_sep_key, algorithm)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_nsec3_hash.py b/tests/test_nsec3_hash.py
index 6f18240..0fd085c 100644
--- a/tests/test_nsec3_hash.py
+++ b/tests/test_nsec3_hash.py
@@ -49,6 +49,15 @@ class NSEC3Hash(unittest.TestCase):
             1,
         ),
         ("*.test-domain.dev", None, 45, "505k9g118d9sofnjhh54rr8fadgpa0ct", 1),
+        (
+            "example",
+            "aabbccdd",
+            12,
+            "0p9mhaveqvm6t7vbl5lop2u3t2rp3tom",
+            dnssec.NSEC3Hash.SHA1
+        ),
+        ("example", "aabbccdd", 12, "0p9mhaveqvm6t7vbl5lop2u3t2rp3tom", "SHA1"),
+        ("example", "aabbccdd", 12, "0p9mhaveqvm6t7vbl5lop2u3t2rp3tom", "sha1")
     ]
 
     def test_hash_function(self):
@@ -67,6 +76,21 @@ class NSEC3Hash(unittest.TestCase):
         with self.assertRaises(ValueError):
             hash = dnssec.nsec3_hash(data[0], data[1], data[2], data[4])
 
+    def test_hash_invalid_algorithm(self):
+        data = (
+            "example.com",
+            "9F1AB450CF71D",
+            0,
+            "qfo2sv6jaej4cm11a3npoorfrckdao2c",
+            1,
+        )
+        with self.assertRaises(ValueError):
+            dnssec.nsec3_hash(data[0], data[1], data[2], 10)
+        with self.assertRaises(ValueError):
+            dnssec.nsec3_hash(data[0], data[1], data[2], "foo")
+
+
+
 
 if __name__ == "__main__":
     unittest.main()