a refactor to the API

3 files changed, 59 insertions, 76 deletions

diff --git a/tests/hazmat/backends/test_multibackend.py b/tests/hazmat/backends/test_multibackend.py
index 57aa7f44..2a533750 100644
--- a/tests/hazmat/backends/test_multibackend.py
+++ b/tests/hazmat/backends/test_multibackend.py
@@ -138,9 +138,8 @@ class DummyCMACBackend(object):
 
 @utils.register_interface(EllipticCurveBackend)
 class DummyEllipticCurveBackend(object):
-    def __init__(self, supported_curves, exchange_supported):
+    def __init__(self, supported_curves):
         self._curves = supported_curves
-        self.exchange_supported = exchange_supported
 
     def elliptic_curve_supported(self, curve):
         return any(
@@ -153,10 +152,7 @@ class DummyEllipticCurveBackend(object):
     ):
         return (
             isinstance(signature_algorithm, ec.ECDSA) and
-            any(
-                isinstance(curve, curve_type)
-                for curve_type in self._curves
-            )
+            self.elliptic_curve_supported(curve)
         )
 
     def generate_elliptic_curve_private_key(self, curve):
@@ -171,8 +167,11 @@ class DummyEllipticCurveBackend(object):
         if not self.elliptic_curve_supported(numbers.curve):
             raise UnsupportedAlgorithm(_Reasons.UNSUPPORTED_ELLIPTIC_CURVE)
 
-    def elliptic_curve_exchange_algorithm_supported(self):
-        return self.exchange_supported
+    def elliptic_curve_exchange_algorithm_supported(self, algorithm, curve):
+        return (
+            isinstance(algorithm, ec.ECDH) and
+            self.elliptic_curve_supported(curve)
+        )
 
 
 @utils.register_interface(PEMSerializationBackend)
@@ -404,7 +403,7 @@ class TestMultiBackend(object):
         backend = MultiBackend([
             DummyEllipticCurveBackend([
                 ec.SECT283K1
-            ], True)
+            ])
         ])
 
         assert backend.elliptic_curve_supported(ec.SECT283K1()) is True
@@ -466,9 +465,13 @@ class TestMultiBackend(object):
                 )
             )
 
-        assert backend.elliptic_curve_exchange_algorithm_supported() is True
-        backend2 = MultiBackend([DummyEllipticCurveBackend([], False)])
-        assert backend2.elliptic_curve_exchange_algorithm_supported() is False
+        assert backend.elliptic_curve_exchange_algorithm_supported(
+            ec.ECDH(), ec.SECT283K1()
+        )
+        backend2 = MultiBackend([DummyEllipticCurveBackend([])])
+        assert not backend2.elliptic_curve_exchange_algorithm_supported(
+            ec.ECDH(), ec.SECT163K1()
+        )
 
     def test_pem_serialization_backend(self):
         backend = MultiBackend([DummyPEMSerializationBackend()])
diff --git a/tests/hazmat/backends/test_openssl.py b/tests/hazmat/backends/test_openssl.py
index 13162046..85331595 100644
--- a/tests/hazmat/backends/test_openssl.py
+++ b/tests/hazmat/backends/test_openssl.py
@@ -534,11 +534,6 @@ class DummyLibrary(object):
     Cryptography_HAS_EC = 0
 
 
-class DummyLibraryECDH(object):
-    Cryptography_HAS_EC = 1
-    Cryptography_HAS_ECDH = 0
-
-
 class TestOpenSSLEllipticCurve(object):
     def test_elliptic_curve_supported(self, monkeypatch):
         monkeypatch.setattr(backend, "_lib", DummyLibrary())
@@ -558,12 +553,9 @@ class TestOpenSSLEllipticCurve(object):
 
     def test_elliptic_curve_exchange_algorithm_supported(self, monkeypatch):
         monkeypatch.setattr(backend, "_lib", DummyLibrary())
-
-        assert backend.elliptic_curve_exchange_algorithm_supported() is False
-
-        monkeypatch.setattr(backend, "_lib", DummyLibraryECDH())
-
-        assert backend.elliptic_curve_exchange_algorithm_supported() is False
+        assert not backend.elliptic_curve_exchange_algorithm_supported(
+            ec.ECDH(), ec.SECP256R1()
+        )
 
 
 @pytest.mark.requires_backend_interface(interface=RSABackend)
diff --git a/tests/hazmat/primitives/test_ec.py b/tests/hazmat/primitives/test_ec.py
index c3a99e5d..2594d5db 100644
--- a/tests/hazmat/primitives/test_ec.py
+++ b/tests/hazmat/primitives/test_ec.py
@@ -57,8 +57,10 @@ def _skip_curve_unsupported(backend, curve):
         )
 
 
-def _skip_exchange_algorithm_unsupported(backend):
-    if not backend.elliptic_curve_exchange_algorithm_supported():
+def _skip_exchange_algorithm_unsupported(backend, algorithm, curve):
+    if not backend.elliptic_curve_exchange_algorithm_supported(
+        algorithm, curve
+    ):
         pytest.skip(
             "Exchange algorithm is not supported by this backend {0}".format(
                 backend
@@ -771,50 +773,6 @@ class DummyECDHBackend(object):
 
 @pytest.mark.requires_backend_interface(interface=EllipticCurveBackend)
 class TestECDHVectors(object):
-
-    def test_unsupported_ecdh_arguments(self, backend):
-        with pytest.raises(TypeError):
-            ec.ECDH(None)
-        curve = ec.SECP521R1
-        _skip_curve_unsupported(backend, curve)
-        prikey = ec.generate_private_key(curve, backend)
-        ecdh = ec.ECDH(prikey)
-        ecdh.compute_key(ecdh.public_key())
-        with pytest.raises(TypeError):
-            ecdh.compute_key(None)
-        with pytest.raises(exceptions.UnsupportedAlgorithm):
-            prikey._backend = DummyECDHBackend()
-            ecdh = ec.ECDH(prikey)
-        _skip_exchange_algorithm_unsupported(DummyECDHBackend())
-
-    def key_exchange(self, backend, vector):
-        key_numbers = vector['IUT']
-        peer_numbers = vector['CAVS']
-
-        prikey = ec.EllipticCurvePrivateNumbers(
-            key_numbers['d'],
-            ec.EllipticCurvePublicNumbers(
-                key_numbers['x'],
-                key_numbers['y'],
-                ec._CURVE_TYPES[vector['curve']]()
-            )
-        ).private_key(backend)
-
-        peerkey = ec.EllipticCurvePrivateNumbers(
-            peer_numbers['d'],
-            ec.EllipticCurvePublicNumbers(
-                peer_numbers['x'],
-                peer_numbers['y'],
-                ec._CURVE_TYPES[vector['curve']]()
-            )
-        ).private_key(backend)
-        peerpubkey = peerkey.public_key()
-
-        ecdh = ec.ECDH(prikey)
-        z = ecdh.compute_key(peerpubkey)
-
-        return int(hexlify(z).decode('ascii'), 16)
-
     @pytest.mark.parametrize(
         "vector",
         load_vectors_from_file(
@@ -825,19 +783,49 @@ class TestECDHVectors(object):
         )
     )
     def test_key_exchange_with_vectors(self, backend, vector):
-        _skip_curve_unsupported(backend, ec._CURVE_TYPES[vector['curve']])
-        _skip_exchange_algorithm_unsupported(backend)
+        _skip_exchange_algorithm_unsupported(
+            backend, ec.ECDH(), ec._CURVE_TYPES[vector['curve']]
+        )
 
+        key_numbers = vector['IUT']
         try:
-            z = self.key_exchange(backend, vector)
+            private_key = ec.EllipticCurvePrivateNumbers(
+                key_numbers['d'],
+                ec.EllipticCurvePublicNumbers(
+                    key_numbers['x'],
+                    key_numbers['y'],
+                    ec._CURVE_TYPES[vector['curve']]()
+                )
+            ).private_key(backend)
         except ValueError:
-            assert vector['fail'] is True
+            # Errno 5 and 6 indicates a bad public key, this doesn't test the
+            # ECDH code at all
+            assert vector['fail'] and vector['errno'] in [5, 6]
+            return
 
-        if vector['fail']:
+        peer_numbers = vector['CAVS']
+        try:
+            peer_pubkey = ec.EllipticCurvePublicNumbers(
+                peer_numbers['x'],
+                peer_numbers['y'],
+                ec._CURVE_TYPES[vector['curve']]()
+            ).public_key(backend)
+        except ValueError:
+            # Errno 1 and 2 indicates a bad public key, this doesn't test the
+            # ECDH code at all
+            assert vector['fail'] and vector['errno'] in [1, 2]
+            return
+
+        if vector['fail'] and vector['errno'] not in [7, 8]:
+            with pytest.raises(ValueError):
+                private_key.exchange(ec.ECDH(), peer_pubkey)
+        else:
+            z = private_key.exchange(ec.ECDH(), peer_pubkey)
+            z = int(hexlify(z).decode('ascii'), 16)
             # Errno 7 denotes a changed private key. Errno 8 denotes a changed
             # shared key. Both these errors will not cause a failure in the
             # exchange but should lead to a non-matching derived shared key.
             if vector['errno'] in [7, 8]:
                 assert z != vector['Z']
-        else:
-            assert z == vector['Z']
+            else:
+                assert z == vector['Z']