Merge pull request #1016 from Ayrx/hkdf-expand-only

HKDF Expand Only implementation

1 files changed, 47 insertions, 17 deletions

diff --git a/cryptography/hazmat/primitives/kdf/hkdf.py b/cryptography/hazmat/primitives/kdf/hkdf.py
index 03500aaa..daa8fcc7 100644
--- a/cryptography/hazmat/primitives/kdf/hkdf.py
+++ b/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -34,6 +34,51 @@ class HKDF(object):
 
         self._algorithm = algorithm
 
+        if isinstance(salt, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as a salt.")
+
+        if salt is None:
+            salt = b"\x00" * (self._algorithm.digest_size // 8)
+
+        self._salt = salt
+
+        self._backend = backend
+
+        self._hkdf_expand = HKDFExpand(self._algorithm, length, info, backend)
+
+    def _extract(self, key_material):
+        h = hmac.HMAC(self._salt, self._algorithm, backend=self._backend)
+        h.update(key_material)
+        return h.finalize()
+
+    def derive(self, key_material):
+        if isinstance(key_material, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as key "
+                "material."
+            )
+
+        return self._hkdf_expand.derive(self._extract(key_material))
+
+    def verify(self, key_material, expected_key):
+        if not constant_time.bytes_eq(self.derive(key_material), expected_key):
+            raise InvalidKey
+
+
+@utils.register_interface(interfaces.KeyDerivationFunction)
+class HKDFExpand(object):
+    def __init__(self, algorithm, length, info, backend):
+        if not isinstance(backend, HMACBackend):
+            raise UnsupportedAlgorithm(
+                "Backend object does not implement HMACBackend",
+                _Reasons.BACKEND_MISSING_INTERFACE
+            )
+
+        self._algorithm = algorithm
+
+        self._backend = backend
+
         max_length = 255 * (algorithm.digest_size // 8)
 
         if length > max_length:
@@ -44,15 +89,6 @@ class HKDF(object):
 
         self._length = length
 
-        if isinstance(salt, six.text_type):
-            raise TypeError(
-                "Unicode-objects must be encoded before using them as a salt.")
-
-        if salt is None:
-            salt = b"\x00" * (self._algorithm.digest_size // 8)
-
-        self._salt = salt
-
         if isinstance(info, six.text_type):
             raise TypeError(
                 "Unicode-objects must be encoded before using them as info.")
@@ -61,15 +97,9 @@ class HKDF(object):
             info = b""
 
         self._info = info
-        self._backend = backend
 
         self._used = False
 
-    def _extract(self, key_material):
-        h = hmac.HMAC(self._salt, self._algorithm, backend=self._backend)
-        h.update(key_material)
-        return h.finalize()
-
     def _expand(self, key_material):
         output = [b""]
         counter = 1
@@ -87,7 +117,7 @@ class HKDF(object):
     def derive(self, key_material):
         if isinstance(key_material, six.text_type):
             raise TypeError(
-                "Unicode-objects must be encoded before using them as key "
+                "Unicode-objects must be encoded before using them as key"
                 "material."
             )
 
@@ -95,7 +125,7 @@ class HKDF(object):
             raise AlreadyFinalized
 
         self._used = True
-        return self._expand(self._extract(key_material))
+        return self._expand(key_material)
 
     def verify(self, key_material, expected_key):
         if not constant_time.bytes_eq(self.derive(key_material), expected_key):