1 files changed, 93 insertions, 1 deletions
diff --git a/tests/hazmat/primitives/test_ec.py b/tests/hazmat/primitives/test_ec.py
index 5467464a..c3a99e5d 100644
--- a/tests/hazmat/primitives/test_ec.py
+++ b/tests/hazmat/primitives/test_ec.py
@@ -7,6 +7,8 @@ from __future__ import absolute_import, division, print_function
 import itertools
 import os
 
+from binascii import hexlify
+
 import pytest
 
 from cryptography import exceptions, utils
@@ -21,7 +23,8 @@ from cryptography.hazmat.primitives.asymmetric.utils import (
 
 from ...utils import (
     load_fips_ecdsa_key_pair_vectors, load_fips_ecdsa_signing_vectors,
-    load_vectors_from_file, raises_unsupported_algorithm
+    load_kasvs_ecdh_vectors, load_vectors_from_file,
+    raises_unsupported_algorithm
 )
 
 _HASH_TYPES = {
@@ -54,6 +57,15 @@ def _skip_curve_unsupported(backend, curve):
         )
 
 
+def _skip_exchange_algorithm_unsupported(backend):
+    if not backend.elliptic_curve_exchange_algorithm_supported():
+        pytest.skip(
+            "Exchange algorithm is not supported by this backend {0}".format(
+                backend
+            )
+        )
+
+
 @utils.register_interface(ec.EllipticCurve)
 class DummyCurve(object):
     name = "dummy-curve"
@@ -749,3 +761,83 @@ class TestECDSAVerification(object):
         public_key = key.public_key()
         with pytest.raises(TypeError):
             public_key.verifier(1234, ec.ECDSA(hashes.SHA256()))
+
+
+class DummyECDHBackend(object):
+    @classmethod
+    def elliptic_curve_exchange_algorithm_supported(cls):
+        return False
+
+
+@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(
+            os.path.join(
+                "asymmetric", "ECDH",
+                "KASValidityTest_ECCStaticUnified_NOKC_ZZOnly_init.fax"),
+            load_kasvs_ecdh_vectors
+        )
+    )
+    def test_key_exchange_with_vectors(self, backend, vector):
+        _skip_curve_unsupported(backend, ec._CURVE_TYPES[vector['curve']])
+        _skip_exchange_algorithm_unsupported(backend)
+
+        try:
+            z = self.key_exchange(backend, vector)
+        except ValueError:
+            assert vector['fail'] is True
+
+        if vector['fail']:
+            # 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']