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

HKDF Expand Only implementation

2 files changed, 65 insertions, 5 deletions

diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
index 2e3c0c3d..598f09f0 100644
--- a/tests/hazmat/primitives/test_hkdf.py
+++ b/tests/hazmat/primitives/test_hkdf.py
@@ -13,6 +13,8 @@
 
 from __future__ import absolute_import, division, print_function
 
+import binascii
+
 import pytest
 
 import six
@@ -21,7 +23,7 @@ from cryptography.exceptions import (
     AlreadyFinalized, InvalidKey, _Reasons
 )
 from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF, HKDFExpand
 
 from ...utils import raises_unsupported_algorithm
 
@@ -151,8 +153,67 @@ class TestHKDF(object):
             hkdf.verify(b"foo", six.u("bar"))
 
 
+@pytest.mark.hmac
+class TestHKDFExpand(object):
+    def test_derive(self, backend):
+        prk = binascii.unhexlify(
+            b"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
+        )
+
+        okm = (b"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c"
+               b"5bf34007208d5b887185865")
+
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info, backend)
+
+        assert binascii.hexlify(hkdf.derive(prk)) == okm
+
+    def test_verify(self, backend):
+        prk = binascii.unhexlify(
+            b"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
+        )
+
+        okm = (b"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c"
+               b"5bf34007208d5b887185865")
+
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info, backend)
+
+        assert hkdf.verify(prk, binascii.unhexlify(okm)) is None
+
+    def test_invalid_verify(self, backend):
+        prk = binascii.unhexlify(
+            b"077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5"
+        )
+
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info, backend)
+
+        with pytest.raises(InvalidKey):
+            hkdf.verify(prk, b"wrong key")
+
+    def test_already_finalized(self, backend):
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info, backend)
+
+        hkdf.derive(b"first")
+
+        with pytest.raises(AlreadyFinalized):
+            hkdf.derive(b"second")
+
+    def test_unicode_error(self, backend):
+        info = binascii.unhexlify(b"f0f1f2f3f4f5f6f7f8f9")
+        hkdf = HKDFExpand(hashes.SHA256(), 42, info, backend)
+
+        with pytest.raises(TypeError):
+            hkdf.derive(six.u("first"))
+
+
 def test_invalid_backend():
     pretend_backend = object()
 
     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..a496459b 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,15 +347,14 @@ 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
     )
 
-    okm = hkdf._expand(binascii.unhexlify(params["prk"]))
+    okm = hkdf.derive(binascii.unhexlify(params["prk"]))
 
     assert okm == binascii.unhexlify(params["okm"])