Add SSL.Context.set_keylog_callback (#910)

* add SSL.Context.set_keylog_callback

* don't fail on missing attribute

* lint!

* make it black

1 files changed, 31 insertions, 0 deletions

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()`)