Add support for setting the SNI hostname header (#139)

* Augment the ssl.wrap_socket method to add support for setting SNI headers

* Update test cases for SNI wrap patch

* fix test

* Switch to ssl.PROTOCOL_SSLv23

* fix doc string to pass pydocstyle test

* fix flake , docs to meet conventions

* fix drop in codecov

2 files changed, 25 insertions, 5 deletions

diff --git a/amqp/transport.py b/amqp/transport.py
index fe0ebe3..e4205ce 100644
--- a/amqp/transport.py
+++ b/amqp/transport.py
@@ -72,7 +72,7 @@ if HAS_TCP_USER_TIMEOUT:
 
 
 try:
-    from socket import TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT # noqa
+    from socket import TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT  # noqa
 except ImportError:
     pass
 else:
@@ -293,13 +293,33 @@ class SSLTransport(_AbstractTransport):
     def _wrap_socket(self, sock, context=None, **sslopts):
         if context:
             return self._wrap_context(sock, sslopts, **context)
-        return ssl.wrap_socket(sock, **sslopts)
+        return self._wrap_socket_sni(sock, **sslopts)
 
     def _wrap_context(self, sock, sslopts, check_hostname=None, **ctx_options):
         ctx = ssl.create_default_context(**ctx_options)
         ctx.check_hostname = check_hostname
         return ctx.wrap_socket(sock, **sslopts)
 
+    def _wrap_socket_sni(sock, keyfile=None, certfile=None,
+                         server_side=False, cert_reqs=ssl.CERT_NONE,
+                         ssl_version=ssl.PROTOCOL_SSLv23, ca_certs=None,
+                         do_handshake_on_connect=True,
+                         suppress_ragged_eofs=True,
+                         server_hostname=None,
+                         ciphers=None):
+        """Socket wrap with SNI headers.
+
+        Default `ssl.wrap_socket` method augmented with support for
+        setting the server_hostname field required for SNI hostname header
+        """
+        sock = ssl.SSLSocket(sock=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,
+                             server_hostname=server_hostname, ciphers=ciphers)
+        return sock
+
     def _shutdown_transport(self):
         """Unwrap a Python 2.6 SSL socket, so we can call shutdown()."""
         if self.sock is not None:
diff --git a/t/unit/test_transport.py b/t/unit/test_transport.py
index 32c4bbd..2eedc33 100644
--- a/t/unit/test_transport.py
+++ b/t/unit/test_transport.py
@@ -323,12 +323,12 @@ class test_SSLTransport:
         self.t.sock.do_handshake.assert_called_with()
         assert self.t._quick_recv is self.t.sock.read
 
-    @patch('ssl.wrap_socket')
-    def test_wrap_socket(self, wrap_socket):
+    def test_wrap_socket(self):
         sock = Mock()
         self.t._wrap_context = Mock()
+        self.t._wrap_socket_sni = Mock()
         self.t._wrap_socket(sock, foo=1)
-        wrap_socket.assert_called_with(sock, foo=1)
+        self.t._wrap_socket_sni.assert_called_with(sock, foo=1)
 
         self.t._wrap_socket(sock, {'c': 2}, foo=1)
         self.t._wrap_context.assert_called_with(sock, {'foo': 1}, c=2)