Fix issue #8806: add SSL contexts support to ftplib

giampaolo · giampaolo · commit a67299e757b6 · 2010-05-26T18:06:04.000Z
diff --git a/Doc/library/ftplib.rst b/Doc/library/ftplib.rst
@@ -65,7 +65,7 @@ The module defines the following items:
       Support for the :keyword:`with` statement was added.
 
 
-.. class:: FTP_TLS(host='', user='', passwd='', acct='', [keyfile[, certfile[, timeout]]])
+.. class:: FTP_TLS(host='', user='', passwd='', acct='', [keyfile[, certfile[, context[, timeout]]]])
 
    A :class:`FTP` subclass which adds TLS support to FTP as described in
    :rfc:`4217`.
@@ -74,6 +74,9 @@ The module defines the following items:
    explicitly ask for it by calling the :meth:`prot_p` method.
    *keyfile* and *certfile* are optional -- they can contain a PEM formatted
    private key and certificate chain file name for the SSL connection.
+   *context* parameter is a :class:`ssl.SSLContext` object which allows
+   bundling SSL configuration options, certificates and private keys into a
+   single (potentially long-lived) structure.
 
    .. versionadded:: 3.2
 
diff --git a/Lib/ftplib.py b/Lib/ftplib.py
@@ -638,9 +638,17 @@ class FTP_TLS(FTP):
         ssl_version = ssl.PROTOCOL_TLSv1
 
         def __init__(self, host='', user='', passwd='', acct='', keyfile=None,
-                     certfile=None, timeout=_GLOBAL_DEFAULT_TIMEOUT):
+                     certfile=None, context=None,
+                     timeout=_GLOBAL_DEFAULT_TIMEOUT):
+            if context is not None and keyfile is not None:
+                raise ValueError("context and keyfile arguments are mutually "
+                                 "exclusive")
+            if context is not None and certfile is not None:
+                raise ValueError("context and certfile arguments are mutually "
+                                 "exclusive")
             self.keyfile = keyfile
             self.certfile = certfile
+            self.context = context
             self._prot_p = False
             FTP.__init__(self, host, user, passwd, acct, timeout)
 
@@ -657,8 +665,12 @@ def auth(self):
                 resp = self.voidcmd('AUTH TLS')
             else:
                 resp = self.voidcmd('AUTH SSL')
-            self.sock = ssl.wrap_socket(self.sock, self.keyfile, self.certfile,
-                                        ssl_version=self.ssl_version)
+            if self.context is not None:
+                self.sock = self.context.wrap_socket(self.sock)
+            else:
+                self.sock = ssl.wrap_socket(self.sock, self.keyfile,
+                                            self.certfile,
+                                            ssl_version=self.ssl_version)
             self.file = self.sock.makefile(mode='r', encoding=self.encoding)
             return resp
 
@@ -689,8 +701,11 @@ def prot_c(self):
         def ntransfercmd(self, cmd, rest=None):
             conn, size = FTP.ntransfercmd(self, cmd, rest)
             if self._prot_p:
-                conn = ssl.wrap_socket(conn, self.keyfile, self.certfile,
-                                       ssl_version=self.ssl_version)
+                if self.context is not None:
+                    conn = self.context.wrap_socket(conn)
+                else:
+                    conn = ssl.wrap_socket(conn, self.keyfile, self.certfile,
+                                           ssl_version=self.ssl_version)
             return conn, size
 
         def retrbinary(self, cmd, callback, blocksize=8192, rest=None):
diff --git a/Lib/test/test_ftplib.py b/Lib/test/test_ftplib.py
@@ -719,6 +719,29 @@ def test_auth_ssl(self):
         finally:
             self.client.ssl_version = ssl.PROTOCOL_TLSv1
 
+    def test_context(self):
+        self.client.quit()
+        ctx = ssl.SSLContext(ssl.PROTOCOL_TLSv1)
+        self.assertRaises(ValueError, ftplib.FTP_TLS, keyfile=CERTFILE,
+                          context=ctx)
+        self.assertRaises(ValueError, ftplib.FTP_TLS, certfile=CERTFILE,
+                          context=ctx)
+        self.assertRaises(ValueError, ftplib.FTP_TLS, certfile=CERTFILE,
+                          keyfile=CERTFILE, context=ctx)
+
+        self.client = ftplib.FTP_TLS(context=ctx, timeout=2)
+        self.client.connect(self.server.host, self.server.port)
+        self.assertNotIsInstance(self.client.sock, ssl.SSLSocket)
+        self.client.auth()
+        self.assertIs(self.client.sock.context, ctx)
+        self.assertIsInstance(self.client.sock, ssl.SSLSocket)
+
+        self.client.prot_p()
+        sock = self.client.transfercmd('list')
+        self.assertIs(self.client.sock.context, ctx)
+        self.assertIsInstance(sock, ssl.SSLSocket)
+        sock.close()
+
 
 class TestTimeouts(TestCase):
 
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -392,6 +392,8 @@ C-API
 Library
 -------
 
+- Issue #8806: add SSL contexts support to ftplib.
+
 - Issue #4769: Fix main() function of the base64 module, use sys.stdin.buffer
   and sys.stdout.buffer (instead of sys.stdin and sys.stdout) to use the bytes
   API
