Add SSL.Context.set_keylog_callback (#910)

* add SSL.Context.set_keylog_callback

* don't fail on missing attribute

* lint!

* make it black

3 files changed, 64 insertions, 1 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 622369d..f7c2c94 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -24,7 +24,8 @@ Deprecations:
 Changes:
 ^^^^^^^^
 
-*none*
+- Added ``Context.set_keylog_callback`` to log key material.
+  `#910 <https://github.com/pyca/pyopenssl/pull/910>`_
 
 
 19.1.0 (2019-11-18)
diff --git a/src/OpenSSL/SSL.py b/src/OpenSSL/SSL.py
index b4b308f..ed20d30 100644
--- a/src/OpenSSL/SSL.py
+++ b/src/OpenSSL/SSL.py
@@ -696,6 +696,11 @@ _requires_alpn = _make_requires(
 )
 
 
+_requires_keylog = _make_requires(
+    getattr(_lib, "Cryptography_HAS_KEYLOG", None), "Key logging not available"
+)
+
+
 class Session(object):
     """
     A class representing an SSL session.  A session defines certain connection
@@ -760,6 +765,7 @@ class Context(object):
         self._verify_helper = None
         self._verify_callback = None
         self._info_callback = None
+        self._keylog_callback = None
         self._tlsext_servername_callback = None
         self._app_data = None
         self._npn_advertise_helper = None
@@ -1338,6 +1344,31 @@ class Context(object):
         )
         _lib.SSL_CTX_set_info_callback(self._context, self._info_callback)
 
+    @_requires_keylog
+    def set_keylog_callback(self, callback):
+        """
+        Set the TLS key logging callback to *callback*. This function will be
+        called whenever TLS key material is generated or received, in order
+        to allow applications to store this keying material for debugging
+        purposes.
+
+        :param callback: The Python callback to use.  This should take two
+            arguments: a Connection object and a bytestring that contains
+            the key material in the format used by NSS for its SSLKEYLOGFILE
+            debugging output.
+        :return: None
+        """
+
+        @wraps(callback)
+        def wrapper(ssl, line):
+            line = _ffi.string(line)
+            callback(Connection._reverse_mapping[ssl], line)
+
+        self._keylog_callback = _ffi.callback(
+            "void (*)(const SSL *, const char *)", wrapper
+        )
+        _lib.SSL_CTX_set_keylog_callback(self._context, self._keylog_callback)
+
     def get_app_data(self):
         """
         Get the application data (supplied via :meth:`set_app_data()`)
diff --git a/tests/test_ssl.py b/tests/test_ssl.py
index ba5b638..a08759f 100644
--- a/tests/test_ssl.py
+++ b/tests/test_ssl.py
@@ -1001,6 +1001,37 @@ class TestContext(object):
             [] == notConnections
         ), "Some info callback arguments were not Connection instances."
 
+    @pytest.mark.skipif(
+        not getattr(_lib, "Cryptography_HAS_KEYLOG", None),
+        reason="SSL_CTX_set_keylog_callback unavailable",
+    )
+    def test_set_keylog_callback(self):
+        """
+        `Context.set_keylog_callback` accepts a callable which will be
+        invoked when key material is generated or received.
+        """
+        called = []
+
+        def keylog(conn, line):
+            called.append((conn, line))
+
+        server_context = Context(TLSv1_METHOD)
+        server_context.set_keylog_callback(keylog)
+        server_context.use_certificate(
+            load_certificate(FILETYPE_PEM, cleartextCertificatePEM)
+        )
+        server_context.use_privatekey(
+            load_privatekey(FILETYPE_PEM, cleartextPrivateKeyPEM)
+        )
+
+        client_context = Context(TLSv1_METHOD)
+
+        self._handshake_test(server_context, client_context)
+
+        assert called
+        assert all(isinstance(conn, Connection) for conn, line in called)
+        assert all(b"CLIENT_RANDOM" in line for conn, line in called)
+
     def _load_verify_locations_test(self, *args):
         """
         Create a client context which will verify the peer certificate and call