From 4df60f18c64ba2835e68bf3eed08d8002a00f4ac Mon Sep 17 00:00:00 2001 From: Christian Heimes Date: Fri, 15 Sep 2017 20:26:05 +0200 Subject: bpo-31386: Custom wrap_bio and wrap_socket type (#3426) SSLSocket.wrap_bio() and SSLSocket.wrap_socket() hard-code SSLObject and SSLSocket as return types. In the light of future deprecation of ssl.wrap_socket() module function and direct instantiation of SSLSocket, it is desirable to make the return type of SSLSocket.wrap_bio() and SSLSocket.wrap_socket() customizable. Signed-off-by: Christian Heimes --- Doc/library/ssl.rst | 36 +++++++++++++++++++--- Lib/ssl.py | 26 +++++++++++----- Lib/test/test_ssl.py | 16 ++++++++++ .../2017-09-07-12-15-56.bpo-27629.7xJXEy.rst | 2 ++ 4 files changed, 67 insertions(+), 13 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2017-09-07-12-15-56.bpo-27629.7xJXEy.rst diff --git a/Doc/library/ssl.rst b/Doc/library/ssl.rst index 200ab04..eb4d8ac 100644 --- a/Doc/library/ssl.rst +++ b/Doc/library/ssl.rst @@ -1593,8 +1593,9 @@ to speed up repeated connections from the same clients. do_handshake_on_connect=True, suppress_ragged_eofs=True, \ server_hostname=None, session=None) - Wrap an existing Python socket *sock* and return an :class:`SSLSocket` - object. *sock* must be a :data:`~socket.SOCK_STREAM` socket; other socket + Wrap an existing Python socket *sock* and return an instance of + :attr:`SSLContext.sslsocket_class` (default :class:`SSLSocket`). + *sock* must be a :data:`~socket.SOCK_STREAM` socket; other socket types are unsupported. The returned SSL socket is tied to the context, its settings and @@ -1617,12 +1618,25 @@ to speed up repeated connections from the same clients. .. versionchanged:: 3.6 *session* argument was added. + .. versionchanged:: 3.7 + The method returns on instance of :attr:`SSLContext.sslsocket_class` + instead of hard-coded :class:`SSLSocket`. + +.. attribute:: SSLContext.sslsocket_class + + The return type of :meth:`SSLContext.wrap_sockets`, defaults to + :class:`SSLSocket`. The attribute can be overridden on instance of class + in order to return a custom subclass of :class:`SSLSocket`. + + .. versionadded:: 3.7 + .. method:: SSLContext.wrap_bio(incoming, outgoing, server_side=False, \ server_hostname=None, session=None) - Create a new :class:`SSLObject` instance by wrapping the BIO objects - *incoming* and *outgoing*. The SSL routines will read input data from the - incoming BIO and write data to the outgoing BIO. + Wrap the BIO objects *incoming* and *outgoing* and return an instance of + attr:`SSLContext.sslobject_class` (default :class:`SSLObject`). The SSL + routines will read input data from the incoming BIO and write data to the + outgoing BIO. The *server_side*, *server_hostname* and *session* parameters have the same meaning as in :meth:`SSLContext.wrap_socket`. @@ -1630,6 +1644,18 @@ to speed up repeated connections from the same clients. .. versionchanged:: 3.6 *session* argument was added. + .. versionchanged:: 3.7 + The method returns on instance of :attr:`SSLContext.sslobject_class` + instead of hard-coded :class:`SSLObject`. + +.. attribute:: SSLContext.sslobject_class + + The return type of :meth:`SSLContext.wrap_bio`, defaults to + :class:`SSLObject`. The attribute can be overridden on instance of class + in order to return a custom subclass of :class:`SSLObject`. + + .. versionadded:: 3.7 + .. method:: SSLContext.session_stats() Get statistics about the SSL sessions created or managed by this context. diff --git a/Lib/ssl.py b/Lib/ssl.py index 062e802..2849dee 100644 --- a/Lib/ssl.py +++ b/Lib/ssl.py @@ -383,10 +383,11 @@ class Purpose(_ASN1Object, _Enum): class SSLContext(_SSLContext): """An SSLContext holds various SSL-related configuration options and data, such as certificates and possibly a private key.""" - - __slots__ = ('protocol', '__weakref__') _windows_cert_stores = ("CA", "ROOT") + sslsocket_class = None # SSLSocket is assigned later. + sslobject_class = None # SSLObject is assigned later. + def __new__(cls, protocol=PROTOCOL_TLS, *args, **kwargs): self = _SSLContext.__new__(cls, protocol) if protocol != _SSLv2_IF_EXISTS: @@ -400,17 +401,21 @@ class SSLContext(_SSLContext): do_handshake_on_connect=True, suppress_ragged_eofs=True, server_hostname=None, session=None): - return SSLSocket(sock=sock, server_side=server_side, - do_handshake_on_connect=do_handshake_on_connect, - suppress_ragged_eofs=suppress_ragged_eofs, - server_hostname=server_hostname, - _context=self, _session=session) + return self.sslsocket_class( + sock=sock, + server_side=server_side, + do_handshake_on_connect=do_handshake_on_connect, + suppress_ragged_eofs=suppress_ragged_eofs, + server_hostname=server_hostname, + _context=self, + _session=session + ) def wrap_bio(self, incoming, outgoing, server_side=False, server_hostname=None, session=None): sslobj = self._wrap_bio(incoming, outgoing, server_side=server_side, server_hostname=server_hostname) - return SSLObject(sslobj, session=session) + return self.sslobject_class(sslobj, session=session) def set_npn_protocols(self, npn_protocols): protos = bytearray() @@ -1135,6 +1140,11 @@ class SSLSocket(socket): return self._sslobj.version() +# Python does not support forward declaration of types. +SSLContext.sslsocket_class = SSLSocket +SSLContext.sslobject_class = SSLObject + + def wrap_socket(sock, keyfile=None, certfile=None, server_side=False, cert_reqs=CERT_NONE, ssl_version=PROTOCOL_TLS, ca_certs=None, diff --git a/Lib/test/test_ssl.py b/Lib/test/test_ssl.py index 523322d..fb5958f 100644 --- a/Lib/test/test_ssl.py +++ b/Lib/test/test_ssl.py @@ -1359,6 +1359,22 @@ class ContextTests(unittest.TestCase): self.assertFalse(ctx.check_hostname) self.assertEqual(ctx.verify_mode, ssl.CERT_NONE) + def test_context_custom_class(self): + class MySSLSocket(ssl.SSLSocket): + pass + + class MySSLObject(ssl.SSLObject): + pass + + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.sslsocket_class = MySSLSocket + ctx.sslobject_class = MySSLObject + + with ctx.wrap_socket(socket.socket(), server_side=True) as sock: + self.assertIsInstance(sock, MySSLSocket) + obj = ctx.wrap_bio(ssl.MemoryBIO(), ssl.MemoryBIO()) + self.assertIsInstance(obj, MySSLObject) + class SSLErrorTests(unittest.TestCase): diff --git a/Misc/NEWS.d/next/Library/2017-09-07-12-15-56.bpo-27629.7xJXEy.rst b/Misc/NEWS.d/next/Library/2017-09-07-12-15-56.bpo-27629.7xJXEy.rst new file mode 100644 index 0000000..95a6c2b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2017-09-07-12-15-56.bpo-27629.7xJXEy.rst @@ -0,0 +1,2 @@ +Make return types of SSLContext.wrap_bio() and SSLContext.wrap_socket() +customizable. -- cgit v0.12