From bcd6ce904d1762e54c4ce31b88fc9c8d91566d1f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jeremy=20Lain=C3=A9?= <jeremy.laine@m4x.org>
Date: Tue, 6 Feb 2018 10:37:53 +0100
Subject: [PATCH] Add Context.set_tlsext_use_srtp

This allows negotiating SRTP keying material, which is useful when using
DTLS-SRTP, as WebRTC does for example.
---
 CHANGELOG.rst      |  3 ++-
 src/OpenSSL/SSL.py | 15 +++++++++++++++
 tests/test_ssl.py  | 29 +++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 244741428..07a515fd0 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -23,7 +23,8 @@ Deprecations:
 Changes:
 ^^^^^^^^
 
-*none*
+- Added ``Context.set_tlsext_use_srtp`` to enable negotiation of SRTP keying material.
+  `#734 <https://github.com/pyca/pyopenssl/pull/734>`_
 
 
 ----
diff --git a/src/OpenSSL/SSL.py b/src/OpenSSL/SSL.py
index 5def0aae8..5c1f92860 100644
--- a/src/OpenSSL/SSL.py
+++ b/src/OpenSSL/SSL.py
@@ -1375,6 +1375,21 @@ def wrapper(ssl, alert, arg):
         _lib.SSL_CTX_set_tlsext_servername_callback(
             self._context, self._tlsext_servername_callback)
 
+    def set_tlsext_use_srtp(self, profiles):
+        """
+        Enable support for negotiating SRTP keying material.
+
+        :param bytes profiles: A colon delimited list of protection profile
+            names, like ``b'SRTP_AES128_CM_SHA1_80:SRTP_AES128_CM_SHA1_32'``.
+        :return: None
+        """
+        if not isinstance(profiles, bytes):
+            raise TypeError("profiles must be a byte string.")
+
+        _openssl_assert(
+            _lib.SSL_CTX_set_tlsext_use_srtp(self._context, profiles) == 0
+        )
+
     @_requires_npn
     def set_npn_advertise_callback(self, callback):
         """
diff --git a/tests/test_ssl.py b/tests/test_ssl.py
index 03dd93524..f57a66fb1 100644
--- a/tests/test_ssl.py
+++ b/tests/test_ssl.py
@@ -1596,6 +1596,35 @@ def test_get_cert_store(self):
         store = context.get_cert_store()
         assert isinstance(store, X509Store)
 
+    def test_set_tlsext_use_srtp_not_bytes(self):
+        """
+        `Context.set_tlsext_use_srtp' enables negotiating SRTP keying material.
+
+        It raises a TypeError if the list of profiles is not a byte string.
+        """
+        context = Context(TLSv1_METHOD)
+        with pytest.raises(TypeError):
+            context.set_tlsext_use_srtp(text_type('SRTP_AES128_CM_SHA1_80'))
+
+    def test_set_tlsext_use_srtp_invalid_profile(self):
+        """
+        `Context.set_tlsext_use_srtp' enables negotiating SRTP keying material.
+
+        It raises an Error if the call to OpenSSL fails.
+        """
+        context = Context(TLSv1_METHOD)
+        with pytest.raises(Error):
+            context.set_tlsext_use_srtp(b'SRTP_BOGUS')
+
+    def test_set_tlsext_use_srtp_valid(self):
+        """
+        `Context.set_tlsext_use_srtp' enables negotiating SRTP keying material.
+
+        It does not return anything.
+        """
+        context = Context(TLSv1_METHOD)
+        assert context.set_tlsext_use_srtp(b'SRTP_AES128_CM_SHA1_80') is None
+
 
 class TestServerNameCallback(object):
     """