Merge branch 'ecdhe' of git://github.com/amluto/pyopenssl into ecdhe

4 files changed, 86 insertions, 1 deletions

diff --git a/.gitignore b/.gitignore
index 867434b..276d4e8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+*.py[co]
 build
 dist
-*.egg-info
\ No newline at end of file
+*.egg-info
diff --git a/OpenSSL/SSL.py b/OpenSSL/SSL.py
index fcc7da4..e6a629b 100644
--- a/OpenSSL/SSL.py
+++ b/OpenSSL/SSL.py
@@ -124,6 +124,18 @@ SSL_CB_CONNECT_EXIT = _lib.SSL_CB_CONNECT_EXIT
 SSL_CB_HANDSHAKE_START = _lib.SSL_CB_HANDSHAKE_START
 SSL_CB_HANDSHAKE_DONE = _lib.SSL_CB_HANDSHAKE_DONE
 
+_Cryptography_HAS_EC = _lib.Cryptography_HAS_EC
+ELLIPTIC_CURVE_DESCRIPTIONS = {}  # In case there's no EC support
+if _Cryptography_HAS_EC:
+    _num_curves = _lib.EC_get_builtin_curves(_ffi.NULL, 0)
+    _curves = _ffi.new('EC_builtin_curve[]', _num_curves)
+    if _lib.EC_get_builtin_curves(_curves, _num_curves) == _num_curves:
+        ELLIPTIC_CURVE_DESCRIPTIONS = dict((_ffi.string(_lib.OBJ_nid2sn(c.nid)),
+                                            _ffi.string(c.comment))
+                                           for c in _curves)
+    del _num_curves
+    del _curves
+
 
 class Error(Exception):
     """
@@ -598,6 +610,35 @@ class Context(object):
         _lib.SSL_CTX_set_tmp_dh(self._context, dh)
 
 
+    def set_tmp_ecdh_curve(self, curve_name):
+        """
+        Select a curve to use for ECDHE key exchange.
+
+        The valid values of *curve_name* are the keys in
+        :py:data:OpenSSL.SSL.ELLIPTIC_CURVE_DESCRIPTIONS.
+
+        Raises a ``ValueError`` if the linked OpenSSL was not compiled with
+        elliptical curve support, or the specified curve is not available.
+
+        :param curve_name: The 'short name' of a curve, e.g. 'prime256v1'
+        :type curve_name: str
+        :return: None
+        """
+        if _lib.Cryptography_HAS_EC:
+            nid = _lib.OBJ_sn2nid(curve_name)
+            if nid == _lib.NID_undef:
+                raise ValueError("No such OpenSSL object '%s'" % curve_name)
+            ecdh = _lib.EC_KEY_new_by_curve_name(nid)
+            if ecdh == _ffi.NULL:
+                raise ValueError(
+                    "OpenSSL could not load the requested elliptic curve"
+                )
+            _lib.SSL_CTX_set_tmp_ecdh(self._context, ecdh)
+            _lib.EC_KEY_free(ecdh)
+        else:
+            raise ValueError("OpenSSL is compiled without ECDH support")
+
+
     def set_cipher_list(self, cipher_list):
         """
         Change the cipher list
diff --git a/OpenSSL/test/test_ssl.py b/OpenSSL/test/test_ssl.py
index a6f0127..0fc8f29 100644
--- a/OpenSSL/test/test_ssl.py
+++ b/OpenSSL/test/test_ssl.py
@@ -35,6 +35,8 @@ from OpenSSL.SSL import (
     SESS_CACHE_OFF, SESS_CACHE_CLIENT, SESS_CACHE_SERVER, SESS_CACHE_BOTH,
     SESS_CACHE_NO_AUTO_CLEAR, SESS_CACHE_NO_INTERNAL_LOOKUP,
     SESS_CACHE_NO_INTERNAL_STORE, SESS_CACHE_NO_INTERNAL)
+from OpenSSL.SSL import (
+    _Cryptography_HAS_EC, ELLIPTIC_CURVE_DESCRIPTIONS)
 
 from OpenSSL.SSL import (
     Error, SysCallError, WantReadError, WantWriteError, ZeroReturnError)
@@ -1172,6 +1174,23 @@ class ContextTests(TestCase, _LoopbackMixin):
         # XXX What should I assert here? -exarkun
 
 
+    def test_set_tmp_ecdh_curve(self):
+        """
+        :py:obj:`Context.set_tmp_ecdh_curve` sets the Eliptical
+        Curve for Diffie-Hellman by the named curve.
+        """
+        context = Context(TLSv1_METHOD)
+        for curve in ELLIPTIC_CURVE_DESCRIPTIONS.keys():
+            context.set_tmp_ecdh_curve(curve)  # Must not throw.
+
+        if _Cryptography_HAS_EC:
+            # If EC is compiled in, there must be at least one curve
+            # Tn theory there could be an OpenSSL that violates this
+            # assumption.  If so, this test will fail and we'll find
+            # out.
+            self.assertTrue(ELLIPTIC_CURVE_DESCRIPTIONS)
+
+
     def test_set_cipher_list_bytes(self):
         """
         :py:obj:`Context.set_cipher_list` accepts a :py:obj:`bytes` naming the
diff --git a/doc/api/ssl.rst b/doc/api/ssl.rst
index da7cfb7..b581450 100644
--- a/doc/api/ssl.rst
+++ b/doc/api/ssl.rst
@@ -110,6 +110,17 @@ Context, Connection.
      .. versionadded:: 0.14
 
 
+.. py:data:: ELLIPTIC_CURVE_DESCRIPTIONS
+
+    A dictionary mapping short names of elliptic curves to textual
+    descriptions.  This dictionary contains exactly the set of curves
+    supported by the OpenSSL build in use.
+
+    The keys are the curve names that can be passed into
+    Constants used with :py:meth:`Context.set_tmp_ecdh_curve` to
+    specify which elliptical curve should be used for ECDHE key exchange.
+
+
 .. py:data:: OPENSSL_VERSION_NUMBER
 
     An integer giving the version number of the OpenSSL library used to build this
@@ -316,6 +327,19 @@ Context objects have the following methods:
 
     Load parameters for Ephemeral Diffie-Hellman from *dhfile*.
 
+.. py:method:: Context.set_tmp_ecdh_curve(curve_name)
+
+   Select a curve to use for ECDHE key exchange.
+
+   The valid values of *curve_name* are the keys in
+   :py:data:`ELLIPTIC_CURVE_DESCRIPTIONS`.
+
+   Raises a ``ValueError`` if the linked OpenSSL was not compiled with
+   elliptical curve support, or the specified curve is not available.
+
+   :param curve_name: The 'short name' of a curve, e.g. 'prime256v1'
+   :type curve_name: str
+   :return: None
 
 .. py:method:: Context.set_app_data(data)