Merge pull request #799 from Ayrx/add-backend-check-to-kdf

Added backend check to kdf primitives

7 files changed, 44 insertions, 8 deletions

diff --git a/cryptography/hazmat/primitives/kdf/hkdf.py b/cryptography/hazmat/primitives/kdf/hkdf.py
index 1a464413..10739178 100644
--- a/cryptography/hazmat/primitives/kdf/hkdf.py
+++ b/cryptography/hazmat/primitives/kdf/hkdf.py
@@ -16,13 +16,20 @@ from __future__ import absolute_import, division, print_function
 import six
 
 from cryptography import utils
-from cryptography.exceptions import AlreadyFinalized, InvalidKey
+from cryptography.exceptions import (
+    AlreadyFinalized, InvalidKey, UnsupportedInterface)
+
+from cryptography.hazmat.backends.interfaces import HMACBackend
 from cryptography.hazmat.primitives import constant_time, hmac, interfaces
 
 
 @utils.register_interface(interfaces.KeyDerivationFunction)
 class HKDF(object):
     def __init__(self, algorithm, length, salt, info, backend):
+        if not isinstance(backend, HMACBackend):
+            raise UnsupportedInterface(
+                "Backend object does not implement HMACBackend")
+
         self._algorithm = algorithm
 
         max_length = 255 * (algorithm.digest_size // 8)
diff --git a/cryptography/hazmat/primitives/kdf/pbkdf2.py b/cryptography/hazmat/primitives/kdf/pbkdf2.py
index 39427780..fcfe601a 100644
--- a/cryptography/hazmat/primitives/kdf/pbkdf2.py
+++ b/cryptography/hazmat/primitives/kdf/pbkdf2.py
@@ -17,14 +17,20 @@ import six
 
 from cryptography import utils
 from cryptography.exceptions import (
-    InvalidKey, UnsupportedHash, AlreadyFinalized
-)
+    InvalidKey, UnsupportedHash, AlreadyFinalized,
+    UnsupportedInterface)
+
+from cryptography.hazmat.backends.interfaces import PBKDF2HMACBackend
 from cryptography.hazmat.primitives import constant_time, interfaces
 
 
 @utils.register_interface(interfaces.KeyDerivationFunction)
 class PBKDF2HMAC(object):
     def __init__(self, algorithm, length, salt, iterations, backend):
+        if not isinstance(backend, PBKDF2HMACBackend):
+            raise UnsupportedInterface(
+                "Backend object does not implement PBKDF2HMACBackend")
+
         if not backend.pbkdf2_hmac_supported(algorithm):
             raise UnsupportedHash(
                 "{0} is not supported for PBKDF2 by this backend".format(
diff --git a/cryptography/hazmat/primitives/twofactor/hotp.py b/cryptography/hazmat/primitives/twofactor/hotp.py
index 27476fd9..34f820c0 100644
--- a/cryptography/hazmat/primitives/twofactor/hotp.py
+++ b/cryptography/hazmat/primitives/twofactor/hotp.py
@@ -25,7 +25,6 @@ from cryptography.hazmat.primitives.hashes import SHA1, SHA256, SHA512
 
 class HOTP(object):
     def __init__(self, key, length, algorithm, backend):
-
         if not isinstance(backend, HMACBackend):
             raise UnsupportedInterface(
                 "Backend object does not implement HMACBackend")
diff --git a/cryptography/hazmat/primitives/twofactor/totp.py b/cryptography/hazmat/primitives/twofactor/totp.py
index 0ce3adaf..08510ef5 100644
--- a/cryptography/hazmat/primitives/twofactor/totp.py
+++ b/cryptography/hazmat/primitives/twofactor/totp.py
@@ -21,7 +21,6 @@ from cryptography.hazmat.primitives.twofactor.hotp import HOTP
 
 class TOTP(object):
     def __init__(self, key, length, algorithm, time_step, backend):
-
         if not isinstance(backend, HMACBackend):
             raise UnsupportedInterface(
                 "Backend object does not implement HMACBackend")
diff --git a/docs/hazmat/primitives/key-derivation-functions.rst b/docs/hazmat/primitives/key-derivation-functions.rst
index 851dbb0b..174b68d2 100644
--- a/docs/hazmat/primitives/key-derivation-functions.rst
+++ b/docs/hazmat/primitives/key-derivation-functions.rst
@@ -84,6 +84,10 @@ Different KDFs are suitable for different tasks such as:
         :class:`~cryptography.hazmat.backends.interfaces.PBKDF2HMACBackend`
         provider.
 
+    :raises cryptography.exceptions.UnsupportedInterface: This is raised if the
+        provided ``backend`` does not implement
+        :class:`~cryptography.hazmat.backends.interfaces.PBKDF2HMACBackend`
+
     .. method:: derive(key_material)
 
         :param bytes key_material: The input key material. For PBKDF2 this
@@ -183,6 +187,10 @@ Different KDFs are suitable for different tasks such as:
         :class:`~cryptography.hazmat.backends.interfaces.HMACBackend`
         provider.
 
+    :raises cryptography.exceptions.UnsupportedInterface: This is raised if the
+        provided ``backend`` does not implement
+        :class:`~cryptography.hazmat.backends.interfaces.HMACBackend`
+
     .. method:: derive(key_material)
 
         :param bytes key_material: The input key material.
diff --git a/tests/hazmat/primitives/test_hkdf.py b/tests/hazmat/primitives/test_hkdf.py
index e3e2a9df..42c75c65 100644
--- a/tests/hazmat/primitives/test_hkdf.py
+++ b/tests/hazmat/primitives/test_hkdf.py
@@ -17,7 +17,9 @@ import six
 
 import pytest
 
-from cryptography.exceptions import AlreadyFinalized, InvalidKey
+from cryptography.exceptions import (
+    AlreadyFinalized, InvalidKey, UnsupportedInterface)
+
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 
@@ -145,3 +147,10 @@ class TestHKDF(object):
             )
 
             hkdf.verify(b"foo", six.u("bar"))
+
+
+def test_invalid_backend():
+    pretend_backend = object()
+
+    with pytest.raises(UnsupportedInterface):
+        HKDF(hashes.SHA256(), 16, None, None, pretend_backend)
diff --git a/tests/hazmat/primitives/test_pbkdf2hmac.py b/tests/hazmat/primitives/test_pbkdf2hmac.py
index f895935b..a47d879e 100644
--- a/tests/hazmat/primitives/test_pbkdf2hmac.py
+++ b/tests/hazmat/primitives/test_pbkdf2hmac.py
@@ -18,8 +18,9 @@ import six
 
 from cryptography import utils
 from cryptography.exceptions import (
-    InvalidKey, UnsupportedHash, AlreadyFinalized
-)
+    InvalidKey, UnsupportedHash, AlreadyFinalized,
+    UnsupportedInterface)
+
 from cryptography.hazmat.primitives import hashes, interfaces
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from cryptography.hazmat.backends import default_backend
@@ -67,3 +68,10 @@ class TestPBKDF2HMAC(object):
         kdf = PBKDF2HMAC(hashes.SHA1(), 20, b"salt", 10, default_backend())
         with pytest.raises(TypeError):
             kdf.derive(six.u("unicode here"))
+
+
+def test_invalid_backend():
+    pretend_backend = object()
+
+    with pytest.raises(UnsupportedInterface):
+        PBKDF2HMAC(hashes.SHA1(), 20, b"salt", 10, pretend_backend)