A class to wrap an SSLSocket in another SSLObject.

SSLWrapper.py

class SSLWrapper:
    """
    SSLWrapper wraps an SSLSocket in an SSLObject.
    This may be used to tunnel an SSL session inside another SSL session.
    
    Example usage:

    sslctx = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH)
    ssock = socket.create_connection(("proxy.example", 443))
    ssock = sslctx.wrap_socket(ssock, server_hostname="localhost")
    wrapper = SSLWrapper()
    wrapper.wrap(sslctx, ssock, server_hostname="server.example")
    wrapper.do_handshake()
    """

    def wrap(self, sslctx: ssl.SSLContext, ssock: ssl.SSLSocket, *args, **kwargs):
        """
        wrap wraps ssock with a new SSLObject.
        The SSLObject is created with sslctx.wrap_bio.
        args and kwargs are forwarded to ssl.wrap_bio.
        """
        self._ssock = ssock
        self._incoming = ssl.MemoryBIO()
        self._outgoing = ssl.MemoryBIO()
        self._sobj = sslctx.wrap_bio(
            self._incoming, self._outgoing, *args, **kwargs)

    def _do(self, fn):
        while True:
            try:
                got = fn()
                return got
            except ssl.SSLWantReadError:
                # Check if there is anything to write.
                if self._outgoing.pending > 0:
                    self._ssock.sendall(self._outgoing.read())
                # Read from underlying socket into the incoming MemoryBIO.
                self._incoming.write(self._ssock.recv())

    def do_handshake(self):
        return self._do(lambda: self._sobj.do_handshake())

    def write(self, data: bytes):
        return self._do(lambda: self._sobj.write(data))

    def read(self, len=1024):
        return self._do(lambda: self._sobj.read(len))