Issue #8808: The IMAP4_SSL constructor now allows passing an SSLContext

parameter to control parameters of the secure channel.  Patch by Sijin
Joseph.

2 files changed, 70 insertions, 8 deletions

diff --git a/Lib/imaplib.py b/Lib/imaplib.py
index 1022e77..142e27b 100644
--- a/Lib/imaplib.py
+++ b/Lib/imaplib.py
@@ -1177,25 +1177,40 @@ if HAVE_SSL:
 
         """IMAP4 client class over SSL connection
 
-        Instantiate with: IMAP4_SSL([host[, port[, keyfile[, certfile]]]])
+        Instantiate with: IMAP4_SSL([host[, port[, keyfile[, certfile[, ssl_context]]]]])
 
                 host - host's name (default: localhost);
-                port - port number (default: standard IMAP4 SSL port).
+                port - port number (default: standard IMAP4 SSL port);
                 keyfile - PEM formatted file that contains your private key (default: None);
                 certfile - PEM formatted certificate chain file (default: None);
+                ssl_context - a SSLContext object that contains your certificate chain
+                              and private key (default: None)
+                Note: if ssl_context is provided, then parameters keyfile or
+                certfile should not be set otherwise ValueError is thrown.
 
         for more documentation see the docstring of the parent class IMAP4.
         """
 
 
-        def __init__(self, host = '', port = IMAP4_SSL_PORT, keyfile = None, certfile = None):
+        def __init__(self, host='', port=IMAP4_SSL_PORT, keyfile=None, certfile=None, ssl_context=None):
+            if ssl_context is not None and keyfile is not None:
+                raise ValueError("ssl_context and keyfile arguments are mutually "
+                                 "exclusive")
+            if ssl_context is not None and certfile is not None:
+                raise ValueError("ssl_context and certfile arguments are mutually "
+                                 "exclusive")
+
             self.keyfile = keyfile
             self.certfile = certfile
+            self.ssl_context = ssl_context
             IMAP4.__init__(self, host, port)
 
         def _create_socket(self):
             sock = IMAP4._create_socket(self)
-            return ssl.wrap_socket(sock, self.keyfile, self.certfile)
+            if self.ssl_context:
+                return self.ssl_context.wrap_socket(sock)
+            else:
+                return ssl.wrap_socket(sock, self.keyfile, self.certfile)
 
         def open(self, host='', port=IMAP4_SSL_PORT):
             """Setup connection to remote server on "host:port".
diff --git a/Lib/test/test_imaplib.py b/Lib/test/test_imaplib.py
index 8034000..c4c7ecc 100644
--- a/Lib/test/test_imaplib.py
+++ b/Lib/test/test_imaplib.py
@@ -258,11 +258,58 @@ class RemoteIMAP_SSLTest(RemoteIMAPTest):
     port = 993
     imap_class = IMAP4_SSL
 
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def create_ssl_context(self):
+        ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+        ssl_context.load_cert_chain(CERTFILE)
+        return ssl_context
+
+    def check_logincapa(self, server):
+        try:
+            for cap in server.capabilities:
+                self.assertIsInstance(cap, str)
+            self.assertFalse('LOGINDISABLED' in server.capabilities)
+            self.assertTrue('AUTH=PLAIN' in server.capabilities)
+            rs = server.login(self.username, self.password)
+            self.assertEqual(rs[0], 'OK')
+        finally:
+            server.logout()
+
     def test_logincapa(self):
-        for cap in self.server.capabilities:
-            self.assertIsInstance(cap, str)
-        self.assertFalse('LOGINDISABLED' in self.server.capabilities)
-        self.assertTrue('AUTH=PLAIN' in self.server.capabilities)
+        with transient_internet(self.host):
+            _server = self.imap_class(self.host, self.port)
+            self.check_logincapa(_server)
+
+    def test_logincapa_with_client_certfile(self):
+        with transient_internet(self.host):
+            _server = self.imap_class(self.host, self.port, certfile=CERTFILE)
+            self.check_logincapa(_server)
+
+    def test_logincapa_with_client_ssl_context(self):
+        with transient_internet(self.host):
+            _server = self.imap_class(self.host, self.port, ssl_context=self.create_ssl_context())
+            self.check_logincapa(_server)
+
+    def test_logout(self):
+        with transient_internet(self.host):
+            _server = self.imap_class(self.host, self.port)
+            rs = _server.logout()
+            self.assertEqual(rs[0], 'BYE')
+
+    def test_ssl_context_certfile_exclusive(self):
+        with transient_internet(self.host):
+            self.assertRaises(ValueError, self.imap_class, self.host, self.port,
+                              certfile=CERTFILE, ssl_context=self.create_ssl_context())
+
+    def test_ssl_context_keyfile_exclusive(self):
+        with transient_internet(self.host):
+            self.assertRaises(ValueError, self.imap_class, self.host, self.port,
+                              keyfile=CERTFILE, ssl_context=self.create_ssl_context())
 
 
 def test_main():