Modified HKDF to use HKDFExpand

1 files changed, 47 insertions, 34 deletions

diff --git a/cryptography/hazmat/primitives/kdf/hkdf.py b/cryptography/hazmat/primitives/kdf/hkdf.py
index 44e14817..d49cc5bd 100644
--- a/cryptography/hazmat/primitives/kdf/hkdf.py
+++ b/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -34,16 +34,6 @@ class HKDF(object):
 
         self._algorithm = algorithm
 
-        max_length = 255 * (algorithm.digest_size // 8)
-
-        if length > max_length:
-            raise ValueError(
-                "Can not derive keys larger than {0} octets.".format(
-                    max_length
-                ))
-
-        self._length = length
-
         if isinstance(salt, six.text_type):
             raise TypeError(
                 "Unicode-objects must be encoded before using them as a salt.")
@@ -53,37 +43,17 @@ class HKDF(object):
 
         self._salt = salt
 
-        if isinstance(info, six.text_type):
-            raise TypeError(
-                "Unicode-objects must be encoded before using them as info.")
-
-        if info is None:
-            info = b""
-
-        self._info = info
         self._backend = backend
 
         self._used = False
 
+        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 _expand(self, key_material):
-        output = [b""]
-        counter = 1
-
-        while (self._algorithm.digest_size // 8) * len(output) < self._length:
-            h = hmac.HMAC(key_material, self._algorithm, backend=self._backend)
-            h.update(output[-1])
-            h.update(self._info)
-            h.update(six.int2byte(counter))
-            output.append(h.finalize())
-            counter += 1
-
-        return b"".join(output)[:self._length]
-
     def derive(self, key_material):
         if isinstance(key_material, six.text_type):
             raise TypeError(
@@ -95,7 +65,7 @@ class HKDF(object):
             raise AlreadyFinalized
 
         self._used = True
-        return self._expand(self._extract(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):
@@ -105,7 +75,50 @@ class HKDF(object):
 @utils.register_interface(interfaces.KeyDerivationFunction)
 class HKDFExpand(HKDF):
     def __init__(self, algorithm, length, info, backend):
-        HKDF.__init__(self, algorithm, length, None, 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:
+            raise ValueError(
+                "Can not derive keys larger than {0} octets.".format(
+                    max_length
+                ))
+
+        self._length = length
+
+        if isinstance(info, six.text_type):
+            raise TypeError(
+                "Unicode-objects must be encoded before using them as info.")
+
+        if info is None:
+            info = b""
+
+        self._info = info
+
+        self._used = False
+
+    def _expand(self, key_material):
+        output = [b""]
+        counter = 1
+
+        while (self._algorithm.digest_size // 8) * len(output) < self._length:
+            h = hmac.HMAC(key_material, self._algorithm, backend=self._backend)
+            h.update(output[-1])
+            h.update(self._info)
+            h.update(six.int2byte(counter))
+            output.append(h.finalize())
+            counter += 1
+
+        return b"".join(output)[:self._length]
 
     def derive(self, key_material):
         if isinstance(key_material, six.text_type):