# Path: autopush/ssl.py
# def monkey_patch_ssl_wrap_socket():
#     """Replace ssl.wrap_socket with ssl_wrap_socket_cached"""
#     ssl.wrap_socket = ssl_wrap_socket_cached
# 
# def ssl_wrap_socket_cached(
#         sock,                          # type: socket.socket
#         keyfile=None,                  # type: Optional[str]
#         certfile=None,                 # type: Optional[str]
#         server_side=False,             # type: bool
#         cert_reqs=ssl.CERT_NONE,       # type: int
#         ssl_version=SSL_PROTO,         # type: int
#         ca_certs=None,                 # type: Optional[str]
#         do_handshake_on_connect=True,  # type: bool
#         suppress_ragged_eofs=True,     # type: bool
#         ciphers=None                   # type: Optional[str]
#         ):
#     # type: (...) -> ssl.SSLSocket
#     """ssl.wrap_socket replacement that caches SSLContexts"""
#     key_kwargs = (
#         ('keyfile', keyfile),
#         ('certfile', certfile),
#         ('cert_reqs', cert_reqs),
#         ('ssl_version', ssl_version),
#         ('ca_certs', ca_certs),
#         ('ciphers', ciphers),
#     )
#     key = frozenset(key_kwargs)
# 
#     context = _sslcontext_cache.get(key)
#     if context is not None:
#         return context.wrap_socket(
#             sock,
#             server_side=server_side,
#             do_handshake_on_connect=do_handshake_on_connect,
#             suppress_ragged_eofs=suppress_ragged_eofs
#         )
# 
#     wrapped = _orig_ssl_wrap_socket(
#         sock,
#         keyfile=keyfile,
#         certfile=certfile,
#         server_side=server_side,
#         cert_reqs=cert_reqs,
#         ssl_version=ssl_version,
#         ca_certs=ca_certs,
#         do_handshake_on_connect=do_handshake_on_connect,
#         suppress_ragged_eofs=suppress_ragged_eofs,
#         ciphers=ciphers
#     )
#     _sslcontext_cache[key] = wrapped.context
#     return wrapped
# 
# def undo_monkey_patch_ssl_wrap_socket():
#     """Undo monkey_patch_ssl_wrap_socket"""
#     ssl.wrap_socket = _orig_ssl_wrap_socket


import socket
import ssl
from twisted.trial import unittest
from autopush.ssl import (
    monkey_patch_ssl_wrap_socket,
    ssl_wrap_socket_cached,
    undo_monkey_patch_ssl_wrap_socket
)



class SSLContextCacheTestCase(unittest.TestCase):

    def setUp(self):
        # XXX: test_main doesn't cleanup after itself
        undo_monkey_patch_ssl_wrap_socket()

    def test_monkey_patch_ssl_wrap_socket(self):
        assert ssl.wrap_socket is not ssl_wrap_socket_cached
        orig = ssl.wrap_socket
        monkey_patch_ssl_wrap_socket()

# Path: autopush/ssl.py
# def monkey_patch_ssl_wrap_socket():
#     """Replace ssl.wrap_socket with ssl_wrap_socket_cached"""
#     ssl.wrap_socket = ssl_wrap_socket_cached
# 
# def ssl_wrap_socket_cached(
#         sock,                          # type: socket.socket
#         keyfile=None,                  # type: Optional[str]
#         certfile=None,                 # type: Optional[str]
#         server_side=False,             # type: bool
#         cert_reqs=ssl.CERT_NONE,       # type: int
#         ssl_version=SSL_PROTO,         # type: int
#         ca_certs=None,                 # type: Optional[str]
#         do_handshake_on_connect=True,  # type: bool
#         suppress_ragged_eofs=True,     # type: bool
#         ciphers=None                   # type: Optional[str]
#         ):
#     # type: (...) -> ssl.SSLSocket
#     """ssl.wrap_socket replacement that caches SSLContexts"""
#     key_kwargs = (
#         ('keyfile', keyfile),
#         ('certfile', certfile),
#         ('cert_reqs', cert_reqs),
#         ('ssl_version', ssl_version),
#         ('ca_certs', ca_certs),
#         ('ciphers', ciphers),
#     )
#     key = frozenset(key_kwargs)
# 
#     context = _sslcontext_cache.get(key)
#     if context is not None:
#         return context.wrap_socket(
#             sock,
#             server_side=server_side,
#             do_handshake_on_connect=do_handshake_on_connect,
#             suppress_ragged_eofs=suppress_ragged_eofs
#         )
# 
#     wrapped = _orig_ssl_wrap_socket(
#         sock,
#         keyfile=keyfile,
#         certfile=certfile,
#         server_side=server_side,
#         cert_reqs=cert_reqs,
#         ssl_version=ssl_version,
#         ca_certs=ca_certs,
#         do_handshake_on_connect=do_handshake_on_connect,
#         suppress_ragged_eofs=suppress_ragged_eofs,
#         ciphers=ciphers
#     )
#     _sslcontext_cache[key] = wrapped.context
#     return wrapped
# 
# def undo_monkey_patch_ssl_wrap_socket():
#     """Undo monkey_patch_ssl_wrap_socket"""
#     ssl.wrap_socket = _orig_ssl_wrap_socket


# Path: autopush/tests/test_ssl.py
import socket
import ssl
from twisted.trial import unittest
from autopush.ssl import (
    monkey_patch_ssl_wrap_socket,
    ssl_wrap_socket_cached,
    undo_monkey_patch_ssl_wrap_socket
)




class SSLContextCacheTestCase(unittest.TestCase):

    def setUp(self):
        # XXX: test_main doesn't cleanup after itself
        undo_monkey_patch_ssl_wrap_socket()

    def test_monkey_patch_ssl_wrap_socket(self):
        assert ssl.wrap_socket is not ssl_wrap_socket_cached
        orig = ssl.wrap_socket
        monkey_patch_ssl_wrap_socket()

        self.addCleanup(undo_monkey_patch_ssl_wrap_socket)