Modified HKDF to use HKDFExpand

3 files changed, 52 insertions, 37 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):
diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
index bee42172..598f09f0 100644
--- a/tests/hazmat/primitives/test_hkdf.py
+++ b/tests/hazmat/primitives/test_hkdf.py
@@ -214,3 +214,6 @@ def test_invalid_backend():
 
     with raises_unsupported_algorithm(_Reasons.BACKEND_MISSING_INTERFACE):
         HKDF(hashes.SHA256(), 16, None, None, pretend_backend)
+
+    with raises_unsupported_algorithm(_Reasons.BACKEND_MISSING_INTERFACE):
+        HKDFExpand(hashes.SHA256(), 16, None, pretend_backend)
diff --git a/tests/hazmat/primitives/utils.py b/tests/hazmat/primitives/utils.py
index 6c3f4c95..7cf5efd0 100644
--- a/tests/hazmat/primitives/utils.py
+++ b/tests/hazmat/primitives/utils.py
@@ -26,7 +26,7 @@ from cryptography.exceptions import (
 from cryptography.hazmat.primitives import hashes, hmac
 from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.hazmat.primitives.ciphers import Cipher
-from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF, HKDFExpand
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 from ...utils import load_vectors_from_file
@@ -347,10 +347,9 @@ def hkdf_extract_test(backend, algorithm, params):
 
 
 def hkdf_expand_test(backend, algorithm, params):
-    hkdf = HKDF(
+    hkdf = HKDFExpand(
         algorithm,
         int(params["l"]),
-        salt=binascii.unhexlify(params["salt"]) or None,
         info=binascii.unhexlify(params["info"]) or None,
         backend=backend
     )