Merge branch 'extra-paramiko-kwargs' of git://github.com/Fizzadar/parallel-ssh into fizzadar

Author: Panos Kittenis

pssh.py

@@ -45,7 +45,7 @@
 DEFAULT_RETRIES = 3
 
 logger = logging.getLogger(__name__)
-    
+
 class UnknownHostException(Exception):
     """Raised when a host is unknown (dns failure)"""
     pass
@@ -74,10 +74,10 @@ class SSHClient(object):
     def __init__(self, host,
                  user=None, password=None, port=None,
                  pkey=None, forward_ssh_agent=True,
-                 num_retries=DEFAULT_RETRIES, _agent=None):
+                 num_retries=DEFAULT_RETRIES, _agent=None, timeout=None):
         """Connect to host honouring any user set configuration in ~/.ssh/config \
         or /etc/ssh/ssh_config
-         
+
         :param host: Hostname to connect to
         :type host: str
         :param user: (Optional) User to login as. Defaults to logged in user or \
@@ -94,6 +94,9 @@ def __init__(self, host,
         :param num_retries: (Optional) Number of retries for connection attempts\
         before the client gives up. Defaults to 3.
         :type num_retries: int
+        :param timeout: (Optional) Number of seconds to timout connection attempts\
+        before the client gives up. Defaults to 10.
+        :type timeout: int
         :param forward_ssh_agent: (Optional) Turn on SSH agent forwarding - \
         equivalent to `ssh -A` from the `ssh` command line utility. \
         Defaults to True if not set.
@@ -140,6 +143,7 @@ def __init__(self, host,
         if _agent:
             self.client._agent = _agent
         self.num_retries = num_retries
+        self.timeout = timeout
         self._connect()
 
     def _connect(self, retries=1):
@@ -148,7 +152,7 @@ def _connect(self, retries=1):
             self.client.connect(self.host, username=self.user,
                                 password=self.password, port=self.port,
                                 pkey=self.pkey,
-                                sock=self.proxy_command)
+                                sock=self.proxy_command, timeout=self.timeout)
         except socket.gaierror, e:
             logger.error("Could not resolve host '%s' - retry %s/%s",
                          self.host, retries, self.num_retries)
@@ -164,8 +168,10 @@ def _connect(self, retries=1):
             while retries < self.num_retries:
                 gevent.sleep(5)
                 return self._connect(retries=retries+1)
+
+            error_type = e.args[1] if len(e.args) > 1 else e.args[0]
             raise ConnectionErrorException("%s for host '%s:%s' - retry %s/%s",
-                                           str(e.args[1]), self.host, self.port,
+                                           str(error_type), self.host, self.port,
                                            retries, self.num_retries,)
         except paramiko.AuthenticationException, e:
             raise AuthenticationException(e)
@@ -229,7 +235,7 @@ def _make_sftp(self):
 
     def mkdir(self, sftp, directory):
         """Make directory via SFTP channel
-        
+
         :param sftp: SFTP client object
         :type sftp: :mod:`paramiko.SFTPClient`
         :param directory: Remote directory to create
@@ -281,7 +287,7 @@ class ParallelSSHClient(object):
 
     def __init__(self, hosts,
                  user=None, password=None, port=None, pkey=None,
-                 forward_ssh_agent=True, num_retries=DEFAULT_RETRIES,
+                 forward_ssh_agent=True, num_retries=DEFAULT_RETRIES, timeout=None,
                  pool_size=10):
         """
         :param hosts: Hosts to connect to
@@ -300,18 +306,21 @@ def __init__(self, hosts,
         :param num_retries: (Optional) Number of retries for connection attempts\
         before the client gives up. Defaults to 3.
         :type num_retries: int
+        :param timeout: (Optional) Number of seconds to timout connection attempts\
+        before the client gives up. Defaults to 10.
+        :type timeout: int
         :param forward_ssh_agent: (Optional) Turn on SSH agent forwarding - \
         equivalent to `ssh -A` from the `ssh` command line utility. \
         Defaults to True if not set.
         :type forward_ssh_agent: bool
         :param pool_size: (Optional) Greenlet pool size. Controls on how many\
         hosts to execute tasks in parallel. Defaults to 10
         :type pool_size: int
-        
+
         **Example**
 
         >>> from pssh import ParallelSSHClient, AuthenticationException,\
-        		UnknownHostException, ConnectionErrorException
+                UnknownHostException, ConnectionErrorException
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'])
         >>> try:
         >>> ... cmds = client.exec_command('ls -ltrh /tmp/aasdfasdf', sudo = True)
@@ -326,7 +335,7 @@ def __init__(self, hosts,
         **Example with returned stdout and stderr buffers**
 
         >>> from pssh import ParallelSSHClient, AuthenticationException,\
-        		UnknownHostException, ConnectionErrorException
+                UnknownHostException, ConnectionErrorException
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'])
         >>> try:
         >>> ... cmds = client.exec_command('ls -ltrh /tmp/aasdfasdf', sudo = True)
@@ -335,10 +344,10 @@ def __init__(self, hosts,
         >>> output = [client.get_stdout(cmd, return_buffers=True) for cmd in cmds]
         >>> print output
         [{'myhost1': {'exit_code': 2,
-        	      'stdout' : <generator object <genexpr>,
+                  'stdout' : <generator object <genexpr>,
                       'stderr' : <generator object <genexpr>,}},
          {'myhost2': {'exit_code': 2,
-         	      'stdout' : <generator object <genexpr>,
+                  'stdout' : <generator object <genexpr>,
                       'stderr' : <generator object <genexpr>,}},
                       ]
         >>> for host_stdout in output:
@@ -352,25 +361,25 @@ def __init__(self, hosts,
         >>> import paramiko
         >>> client_key = paramiko.RSAKey.from_private_key_file('user.key')
         >>> client = ParallelSSHClient(['myhost1', 'myhost2'], pkey=client_key)
-        
+
         .. note ::
-          
+
           **Connection persistence**
-          
+
           Connections to hosts will remain established for the duration of the
           object's life. To close them, just `del` or reuse the object reference.
-          
+
           >>> client = ParallelSSHClient(['localhost'])
           >>> cmds = client.exec_command('ls -ltrh /tmp/aasdfasdf')
           >>> cmds[0].join()
-          
+
           :netstat: ``tcp        0      0 127.0.0.1:53054         127.0.0.1:22            ESTABLISHED``
-          
+
           Connection remains active after commands have finished executing. Any \
           additional commands will use the same connection.
-          
+
           >>> del client
-          
+
           Connection is terminated.
         """
         self.pool = gevent.pool.Pool(size=pool_size)
@@ -382,6 +391,7 @@ def __init__(self, hosts,
         self.port = port
         self.pkey = pkey
         self.num_retries = num_retries
+        self.timeout = timeout
         # To hold host clients
         self.host_clients = dict((host, None) for host in hosts)
 
@@ -396,39 +406,40 @@ def exec_command(self, *args, **kwargs):
         :rtype: List of :mod:`gevent.Greenlet`
 
         **Example**:
-      
+
         >>> cmds = client.exec_command('ls -ltrh')
-        
+
         Wait for completion, no stdout:
-        
+
         >>> for cmd in cmds:
         >>>     cmd.join()
-        
+
         Alternatively/in addition print stdout for each command:
-        
+
         >>> print [get_stdout(cmd) for cmd in cmds]
 
         Retrieving stdout implies join, meaning get_stdout will wait
         for completion of all commands before returning output.
-        
+
         You may call get_stdout on already completed greenlets to re-get
         their output as many times as you want."""
         return [self.pool.spawn(self._exec_command, host, *args, **kwargs)
                 for host in self.hosts]
-    
+
     def _exec_command(self, host, *args, **kwargs):
         """Make SSHClient, run command on host"""
         if not self.host_clients[host]:
             self.host_clients[host] = SSHClient(host, user=self.user,
                                                 password=self.password,
                                                 port=self.port, pkey=self.pkey,
                                                 forward_ssh_agent=self.forward_ssh_agent,
-                                                num_retries=self.num_retries)
+                                                num_retries=self.num_retries,
+                                                timeout=self.timeout)
         return self.host_clients[host].exec_command(*args, **kwargs)
 
     def get_stdout(self, greenlet, return_buffers=False):
         """Get/print stdout from greenlet and return exit code for host
-        
+
         :mod:`pssh.get_stdout` will close the open SSH channel but this does
         **not** close the established connection to the remote host, only the
         authenticated SSH channel within it. This is standard practise
@@ -452,7 +463,7 @@ def get_stdout(self, greenlet, return_buffers=False):
         for example ``{'myhost1': {'exit_code': 0}}``
         :rtype: With ``return_buffers=True``: ``{'myhost1': {'exit_code': 0,
                                                              'channel' : None or SSH channel of command if command is still executing,
-         						     'stdout' : <iterable>,
+                                     'stdout' : <iterable>,
                                                              'stderr' : <iterable>,}}``
         """
         gevent.sleep(.2)
@@ -488,7 +499,7 @@ def get_stdout(self, greenlet, return_buffers=False):
 
     def copy_file(self, local_file, remote_file):
         """Copy local file to remote file in parallel
-        
+
         :param local_file: Local filepath to copy to remote host
         :type local_file: str
         :param remote_file: Remote filepath on remote host to copy file to

tests/test_ssh_client.py

@@ -35,7 +35,7 @@
     os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key']))
 
 class SSHClientTest(unittest.TestCase):
-    
+
     def setUp(self):
         self.fake_cmd = 'fake cmd'
         self.fake_resp = 'fake response'
@@ -133,6 +133,13 @@ def test_ssh_client_conn_failure(self):
                           SSHClient, host, port=self.listen_port,
                           pkey=self.user_key, num_retries=0)
 
+    def test_ssh_client_timeout(self):
+        """Test connection timeout error"""
+        with self.assertRaises(ConnectionErrorException) as cm:
+            SSHClient('127.0.0.1', port=self.listen_port,
+                      pkey=self.user_key, num_retries=0, timeout=1)
+
+        self.assertEqual(cm.exception.args[1], 'timed out')
 
 if __name__ == '__main__':
     unittest.main()