Break up hkdf_derive into hkdf_extract and hkdf_expand.

Testing each individually against all the vectors and actually
asserting about the intermediate state.

hkdf_derive is now just a helper function which copes with the
default arguments.

1 files changed, 24 insertions, 23 deletions

diff --git a/cryptography/hazmat/primitives/kdf/hkdf.py b/cryptography/hazmat/primitives/kdf/hkdf.py
index 8d36c80b..3f3897c1 100644
--- a/cryptography/hazmat/primitives/kdf/hkdf.py
+++ b/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -11,42 +11,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import six
+
 from cryptography.hazmat.primitives import hmac
-from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives import constant_time
 
-def hkdf_derive(input_key, key_length, salt, info, hash, backend):
-    if hash is None:
-        hash = hashes.SHA256()
 
-    if info is None:
-        info = b""
-
-    if salt is None:
-        salt = b"\x00" * (hash.digest_size // 8)
+def hkdf_extract(algorithm, ikm, salt, backend):
+    h = hmac.HMAC(salt, algorithm, backend=backend)
+    h.update(ikm)
+    return h.finalize()
 
-    h = hmac.HMAC(salt, hash, backend=backend)
-    h.update(input_key)
-    PRK = h.finalize()
 
+def hkdf_expand(algorithm, prk, info, length, backend):
     output = [b'']
     counter = 1
 
-    while (hash.digest_size // 8) * len(output) < key_length:
-        h = hmac.HMAC(PRK, hash, backend=backend)
+    while (algorithm.digest_size // 8) * len(output) < length:
+        h = hmac.HMAC(prk, algorithm, backend=backend)
         h.update(output[-1])
         h.update(info)
-        h.update(chr(counter))
+        h.update(six.int2byte(counter))
         output.append(h.finalize())
         counter += 1
 
-    return b"".join(output)[:key_length]
+    return b"".join(output)[:length]
 
 
-def hkdf_verify(expected, input_key, key_length, salt, info, hash, backend):
-    derived = hkdf_derive(input_key, key_length, salt=salt, info=info,
-                           hash=hash, backend=backend)
-
-    if not constant_time.bytes_eq(expected, derived):
-        raise ValueError("")
+def hkdf_derive(key, length, salt, info, algorithm, backend):
+    if info is None:
+        info = b""
 
+    if salt is None:
+        salt = b"\x00" * (algorithm.digest_size // 8)
+
+    return hkdf_expand(
+        algorithm,
+        hkdf_extract(algorithm, key, salt, backend=backend),
+        info,
+        length,
+        backend=backend
+    )