From feb134586ee6ca56e2c53b35d0ffbb79eb1b5dee Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Wed, 18 Jul 2018 05:20:48 -0500
Subject: [PATCH] Refs #3331 -- integrated wycheproof ECDH tests (#4354)

* Refs #3331 -- integrated wycheproof ECDH tests

* flake8 + missing assert

* Handle this error case

* skip on unsupported

* shouldn't need to try here any more
---
 tests/utils.py                |  4 ++
 tests/wycheproof/test_ecdh.py | 83 +++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 tests/wycheproof/test_ecdh.py

diff --git a/tests/utils.py b/tests/utils.py
index ccc3b7c1bbea..b950f8bd9996 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -906,6 +906,10 @@ def valid(self):
     def acceptable(self):
         return self.testcase["result"] == "acceptable"
 
+    @property
+    def invalid(self):
+        return self.testcase["result"] == "invalid"
+
     def has_flag(self, flag):
         return flag in self.testcase["flags"]
 
diff --git a/tests/wycheproof/test_ecdh.py b/tests/wycheproof/test_ecdh.py
new file mode 100644
index 000000000000..0850b627ddbd
--- /dev/null
+++ b/tests/wycheproof/test_ecdh.py
@@ -0,0 +1,83 @@
+# This file is dual licensed under the terms of the Apache License, Version
+# 2.0, and the BSD License. See the LICENSE file in the root of this repository
+# for complete details.
+
+from __future__ import absolute_import, division, print_function
+
+import binascii
+
+import pytest
+
+from cryptography.exceptions import UnsupportedAlgorithm
+from cryptography.hazmat.backends.interfaces import EllipticCurveBackend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import ec
+
+from ..hazmat.primitives.test_ec import _skip_exchange_algorithm_unsupported
+
+
+_CURVES = {
+    "secp224r1": ec.SECP224R1(),
+    "secp256r1": ec.SECP256R1(),
+    "secp384r1": ec.SECP384R1(),
+    "secp521r1": ec.SECP521R1(),
+    "secp256k1": ec.SECP256K1(),
+    "brainpoolP224r1": None,
+    "brainpoolP256r1": ec.BrainpoolP256R1(),
+    "brainpoolP320r1": None,
+    "brainpoolP384r1": ec.BrainpoolP384R1(),
+    "brainpoolP512r1": ec.BrainpoolP512R1(),
+    "brainpoolP224t1": None,
+    "brainpoolP256t1": None,
+    "brainpoolP320t1": None,
+    "brainpoolP384t1": None,
+    "brainpoolP512t1": None,
+}
+
+
+@pytest.mark.requires_backend_interface(interface=EllipticCurveBackend)
+@pytest.mark.wycheproof_tests(
+    "ecdh_test.json",
+    "ecdh_brainpoolP224r1_test.json",
+    "ecdh_brainpoolP256r1_test.json",
+    "ecdh_brainpoolP320r1_test.json",
+    "ecdh_brainpoolP384r1_test.json",
+    "ecdh_brainpoolP512r1_test.json",
+    "ecdh_secp224r1_test.json",
+    "ecdh_secp256k1_test.json",
+    "ecdh_secp256r1_test.json",
+    "ecdh_secp384r1_test.json",
+    "ecdh_secp521r1_test.json",
+)
+def test_ecdh(backend, wycheproof):
+    curve = _CURVES[wycheproof.testcase["curve"]]
+    if curve is None:
+        pytest.skip(
+            "Unsupported curve ({})".format(wycheproof.testcase["curve"])
+        )
+    _skip_exchange_algorithm_unsupported(backend, ec.ECDH(), curve)
+
+    private_key = ec.derive_private_key(
+        int(wycheproof.testcase["private"], 16), curve, backend
+    )
+
+    try:
+        public_key = serialization.load_der_public_key(
+            binascii.unhexlify(wycheproof.testcase["public"]), backend
+        )
+    except NotImplementedError:
+        assert wycheproof.has_flag("UnnamedCurve")
+        return
+    except ValueError:
+        assert wycheproof.invalid or wycheproof.acceptable
+        return
+    except UnsupportedAlgorithm:
+        return
+
+    if wycheproof.valid or wycheproof.acceptable:
+        computed_shared = private_key.exchange(ec.ECDH(), public_key)
+        expected_shared = binascii.unhexlify(wycheproof.testcase["shared"])
+        assert computed_shared == expected_shared
+    else:
+        with pytest.raises(ValueError):
+            private_key.exchange(ec.ECDH(), public_key)