Nice error message on unsupported hash function

data61 · May 10, 2018 · db8d106 · db8d106
1 parent 3985ea2
commit db8d106
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/clkhash/key_derivation.py b/clkhash/key_derivation.py
@@ -4,6 +4,7 @@
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 from cryptography.hazmat.backends import default_backend
 from future.builtins import range, zip
+from future.utils import raise_from
 
 
 """
@@ -64,7 +65,14 @@ def hkdf(master_secret,             # type: bytes
     :param key_size: the size of the produced keys
     :return: Derived keys
     """
-    hkdf = HKDF(algorithm=_HASH_FUNCTIONS[hash_algo](),
+    try:
+        hash_function = _HASH_FUNCTIONS[hash_algo]
+    except KeyError:
+        msg = "unsupported hash function '{}'".format(hash_algo)
+        raise_from(ValueError(msg), None)
+
+
+    hkdf = HKDF(algorithm=hash_function(),
                 length=num_keys * key_size,
                 salt=salt,
                 info=info,

diff --git a/tests/test_key_derivation.py b/tests/test_key_derivation.py
@@ -129,3 +129,8 @@ def test_wrong_kdf(self):
         with self.assertRaises(ValueError):
             generate_key_lists([b'0'], 1, kdf='breakMe')
 
+    def test_wrong_hash_function(self):
+        with self.assertRaises(ValueError):
+            hkdf('foo'.encode('ascii'),
+                 3,
+                 hash_algo='obviously_unsupported')