Fix issue 8594: adds a source_address parameter to ftplib module.

giampaolo · giampaolo · commit 396ff060519a · 2011-02-28T19:19:51.000Z
diff --git a/Doc/library/ftplib.rst b/Doc/library/ftplib.rst
@@ -40,15 +40,16 @@ Here's a sample session using the :mod:`ftplib` module::
 
 The module defines the following items:
 
-.. class:: FTP(host='', user='', passwd='', acct=''[, timeout])
+.. class:: FTP(host='', user='', passwd='', acct='', timeout=None, source_address=None)
 
    Return a new instance of the :class:`FTP` class.  When *host* is given, the
    method call ``connect(host)`` is made.  When *user* is given, additionally
    the method call ``login(user, passwd, acct)`` is made (where *passwd* and
    *acct* default to the empty string when not given).  The optional *timeout*
    parameter specifies a timeout in seconds for blocking operations like the
    connection attempt (if is not specified, the global default timeout setting
-   will be used).
+   will be used). *source_address* is a 2-tuple ``(host, port)`` for the socket
+   to bind to as its source address before connecting.
 
    :class:`FTP` class supports the :keyword:`with` statement. Here is a sample
    on how using it:
@@ -68,8 +69,11 @@ The module defines the following items:
    .. versionchanged:: 3.2
       Support for the :keyword:`with` statement was added.
 
+   .. versionchanged:: 3.3
+      *source_address* parameter was added.
 
-.. class:: FTP_TLS(host='', user='', passwd='', acct='', [keyfile[, certfile[, context[, timeout]]]])
+
+.. class:: FTP_TLS(host='', user='', passwd='', acct='', keyfile=None, certfile=None, context=None, timeout=None, source_address=None)
 
    A :class:`FTP` subclass which adds TLS support to FTP as described in
    :rfc:`4217`.
@@ -80,10 +84,15 @@ The module defines the following items:
    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.
+   single (potentially long-lived) structure. *source_address* is a 2-tuple
+   ``(host, port)`` for the socket to bind to as its source address before
+   connecting.
 
    .. versionadded:: 3.2
 
+   .. versionchanged:: 3.3
+      *source_address* parameter was added.
+
    Here's a sample session using the :class:`FTP_TLS` class:
 
    >>> from ftplib import FTP_TLS
@@ -174,18 +183,22 @@ followed by ``lines`` for the text version or ``binary`` for the binary version.
    debugging output, logging each line sent and received on the control connection.
 
 
-.. method:: FTP.connect(host='', port=0[, timeout])
+.. method:: FTP.connect(host='', port=0, timeout=None, source_address=None)
 
    Connect to the given host and port.  The default port number is ``21``, as
    specified by the FTP protocol specification.  It is rarely needed to specify a
    different port number.  This function should be called only once for each
    instance; it should not be called at all if a host was given when the instance
    was created.  All other methods can only be used after a connection has been
    made.
-
    The optional *timeout* parameter specifies a timeout in seconds for the
    connection attempt. If no *timeout* is passed, the global default timeout
    setting will be used.
+   *source_address* is a 2-tuple ``(host, port)`` for the socket to bind to as
+   its source address before connecting.
+
+   .. versionchanged:: 3.3
+      *source_address* parameter was added.
 
 
 .. method:: FTP.getwelcome()
diff --git a/Lib/ftplib.py b/Lib/ftplib.py
@@ -107,7 +107,8 @@ class FTP:
     # Optional arguments are host (for connect()),
     # and user, passwd, acct (for login())
     def __init__(self, host='', user='', passwd='', acct='',
-                 timeout=_GLOBAL_DEFAULT_TIMEOUT):
+                 timeout=_GLOBAL_DEFAULT_TIMEOUT, source_address=None):
+        self.source_address = source_address
         self.timeout = timeout
         if host:
             self.connect(host)
@@ -128,18 +129,23 @@ def __exit__(self, *args):
                 if self.sock is not None:
                     self.close()
 
-    def connect(self, host='', port=0, timeout=-999):
+    def connect(self, host='', port=0, timeout=-999, source_address=None):
         '''Connect to host.  Arguments are:
          - host: hostname to connect to (string, default previous host)
          - port: port to connect to (integer, default previous port)
+         - source_address: a 2-tuple (host, port) for the socket to bind
+           to as its source address before connecting.
         '''
         if host != '':
             self.host = host
         if port > 0:
             self.port = port
         if timeout != -999:
             self.timeout = timeout
-        self.sock = socket.create_connection((self.host, self.port), self.timeout)
+        if source_address is not None:
+            self.source_address = source_address
+        self.sock = socket.create_connection((self.host, self.port), self.timeout,
+                                             source_address=self.source_address)
         self.af = self.sock.family
         self.file = self.sock.makefile('r', encoding=self.encoding)
         self.welcome = self.getresp()
@@ -334,7 +340,8 @@ def ntransfercmd(self, cmd, rest=None):
         size = None
         if self.passiveserver:
             host, port = self.makepasv()
-            conn = socket.create_connection((host, port), self.timeout)
+            conn = socket.create_connection((host, port), self.timeout,
+                                            source_address=self.source_address)
             if rest is not None:
                 self.sendcmd("REST %s" % rest)
             resp = self.sendcmd(cmd)
@@ -637,7 +644,7 @@ class FTP_TLS(FTP):
 
         def __init__(self, host='', user='', passwd='', acct='', keyfile=None,
                      certfile=None, context=None,
-                     timeout=_GLOBAL_DEFAULT_TIMEOUT):
+                     timeout=_GLOBAL_DEFAULT_TIMEOUT, source_address=None):
             if context is not None and keyfile is not None:
                 raise ValueError("context and keyfile arguments are mutually "
                                  "exclusive")
@@ -648,7 +655,7 @@ def __init__(self, host='', user='', passwd='', acct='', keyfile=None,
             self.certfile = certfile
             self.context = context
             self._prot_p = False
-            FTP.__init__(self, host, user, passwd, acct, timeout)
+            FTP.__init__(self, host, user, passwd, acct, timeout, source_address)
 
         def login(self, user='', passwd='', acct='', secure=True):
             if secure and not isinstance(self.sock, ssl.SSLSocket):
diff --git a/Lib/test/test_ftplib.py b/Lib/test/test_ftplib.py
@@ -608,6 +608,20 @@ def is_client_connected():
         self.assertEqual(self.server.handler_instance.last_received_cmd, 'quit')
         self.assertFalse(is_client_connected())
 
+    def test_source_address(self):
+        self.client.quit()
+        port = support.find_unused_port()
+        self.client.connect(self.server.host, self.server.port,
+                            source_address=(HOST, port))
+        self.assertEqual(self.client.sock.getsockname()[1], port)
+        self.client.quit()
+
+    def test_source_address_passive_connection(self):
+        port = support.find_unused_port()
+        self.client.source_address = (HOST, port)
+        sock = self.client.transfercmd('list')
+        self.assertEqual(sock.getsockname()[1], port)
+
     def test_parse257(self):
         self.assertEqual(ftplib.parse257('257 "/foo/bar"'), '/foo/bar')
         self.assertEqual(ftplib.parse257('257 "/foo/bar" created'), '/foo/bar')
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -35,6 +35,9 @@ Core and Builtins
 Library
 -------
 
+- Issue 8594: ftplib now provides a source_address parameter to specify which
+  (address, port) to bind to before connecting.
+
 - Issue #11326: Add the missing connect_ex() implementation for SSL sockets,
   and make it work for non-blocking connects.
 
