Cleanup. Updated docstrings. Fixed missing pty regression. Added pty test for ssh client. Updated timeout test to check requested timeout against timeout set in channel.

Author: Dan

pssh.py

@@ -147,16 +147,16 @@ def __init__(self, host,
         if self.proxy_host and self.proxy_port:
             logger.debug("Proxy configured for destination host %s - Proxy host: %s:%s",
                          self.host, self.proxy_host, self.proxy_port,)
-            self._connect_proxy()
+            self._connect_tunnel()
         else:
             self._connect(self.client, self.host, self.port)
 
-    def _connect_proxy(self):
-        """Connects to SSH server and returns transport channel to be used
-        to forward proxy to another SSH server.
-        client (me) -> proxy (ssh server to proxy through) -> \
+    def _connect_tunnel(self):
+        """Connects to SSH server via an intermediate SSH tunnel server.
+        client (me) -> tunnel (ssh server to proxy through) -> \
         destination (ssh server to run command)
-        :rtype: `:mod:paramiko.Channel` Channel to proxy through"""
+        :rtype: `:mod:paramiko.SSHClient` Client to remote SSH destination
+        via intermediate SSH tunnel server."""
         self.proxy_client = paramiko.SSHClient()
         self.proxy_client.set_missing_host_key_policy(paramiko.MissingHostKeyPolicy())
         self._connect(self.proxy_client, self.proxy_host, self.proxy_port)
@@ -168,44 +168,51 @@ def _connect_proxy(self):
         return self._connect(self.client, self.host, self.port, sock=proxy_channel)
 
     def _connect(self, client, host, port, sock=None, retries=1):
-        """Connect to host, throw UnknownHost exception on DNS errors"""
+        """Connect to host
+        
+        :raises: :mod:`pssh.AuthenticationException` on authentication error
+        :raises: :mod:`pssh.UnknownHostException` on DNS resolution error
+        :raises: :mod:`pssh.ConnectionErrorException` on error connecting
+        :raises: :mod:`pssh.SSHException` on other undefined SSH errors
+        """
         try:
             client.connect(host, username=self.user,
                            password=self.password, port=port,
                            pkey=self.pkey,
                            sock=sock, timeout=self.timeout)
-        except sock_gaierror, e:
+        except sock_gaierror, ex:
             logger.error("Could not resolve host '%s' - retry %s/%s",
                          self.host, retries, self.num_retries)
             while retries < self.num_retries:
                 gevent.sleep(5)
                 return self._connect(client, host, port, sock=sock,
                                      retries=retries+1)
             raise UnknownHostException("%s - %s - retry %s/%s",
-                                       str(e.args[1]),
+                                       str(ex.args[1]),
                                        self.host, retries, self.num_retries)
-        except sock_error, e:
+        except sock_error, ex:
             logger.error("Error connecting to host '%s:%s' - retry %s/%s",
                          self.host, self.port, retries, self.num_retries)
             while retries < self.num_retries:
                 gevent.sleep(5)
                 return self._connect(client, host, port, sock=sock,
                                      retries=retries+1)
-            error_type = e.args[1] if len(e.args) > 1 else e.args[0]
+            error_type = ex.args[1] if len(ex.args) > 1 else ex.args[0]
             raise ConnectionErrorException("%s for host '%s:%s' - retry %s/%s",
                                            str(error_type), self.host, self.port,
                                            retries, self.num_retries,)
-        except paramiko.AuthenticationException, e:
-            raise AuthenticationException(e)
+        except paramiko.AuthenticationException, ex:
+            raise AuthenticationException(ex)
         # SSHException is more general so should be below other types
         # of SSH failure
-        except paramiko.SSHException, e:
-            raise SSHException(e)
+        except paramiko.SSHException, ex:
+            logger.error("General SSH error - %s", ex)
+            raise SSHException(ex)
 
     def exec_command(self, command, sudo=False, user=None, **kwargs):
         """Wrapper to :mod:`paramiko.SSHClient.exec_command`
 
-        Opens a new SSH session with a pty and runs command with given \
+        Opens a new SSH session with a new pty and runs command with given \
         `kwargs` if any. Greenlet then yields (sleeps) while waiting for \
         command to finish executing or channel to close indicating the same.
 
@@ -224,7 +231,7 @@ def exec_command(self, command, sudo=False, user=None, **kwargs):
         channel = self.client.get_transport().open_session()
         if self.forward_ssh_agent:
             agent_handler = paramiko.agent.AgentRequestHandler(channel)
-        # channel.get_pty()
+        channel.get_pty()
         if self.timeout:
             channel.settimeout(self.timeout)
         _stdout, _stderr = channel.makefile('rb'), \

tests/test_pssh_client.py

@@ -191,27 +191,32 @@ def test_pssh_client_auth_failure(self):
         server.join()
 
     def test_pssh_client_timeout(self):
+        server_timeout=0.2
+        client_timeout=server_timeout-0.1
         server = start_server({ self.fake_cmd : self.fake_resp },
                               self.listen_socket,
-                              timeout=0.2)
+                              timeout=server_timeout)
         client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
                                    pkey=self.user_key,
-                                   timeout=0.1)
-        cmd = client.exec_command(self.fake_cmd)[0]
+                                   timeout=client_timeout)
+        output = client.run_command(self.fake_cmd)
         # Handle exception
         try:
-            gevent.sleep(0.5)
-            cmd.get()
+            gevent.sleep(server_timeout+0.2)
+            client.pool.join()
             if not server.exception:
                 raise Exception(
                     "Expected gevent.Timeout from socket timeout, got none")
             raise server.exception
         except gevent.Timeout:
             pass
+        chan_timeout = output['127.0.0.1']['channel'].gettimeout()
+        self.assertEqual(client_timeout, chan_timeout,
+                         msg="Channel timeout %s does not match requested timeout %s" %(
+                             chan_timeout, client_timeout,))
         del client
         server.join()
 
-
     def test_pssh_client_exec_command_password(self):
         """Test password authentication. Fake server accepts any password
         even empty string"""

tests/test_ssh_client.py

@@ -136,6 +136,21 @@ def test_ssh_client_unknown_host_failure(self):
                           SSHClient, host, port=self.listen_port,
                           pkey=self.user_key, num_retries=0)
 
+    def test_ssh_client_pty(self):
+        """Test that we get a new pty for our non-interactive SSH sessions"""
+        server = start_server({ self.fake_cmd : self.fake_resp },
+                                self.listen_socket)
+        client = SSHClient('127.0.0.1', port=self.listen_port)
+        channel = client.client.get_transport().open_session()
+        self.assertFalse(channel.event.is_set(),
+                         msg="Got pty without requesting it")
+        channel.get_pty()
+        self.assertTrue(channel.event.is_set(),
+                        msg="Requested pty but got none")
+        channel.close()
+        del channel
+        del client
+        server.join()
 
 if __name__ == '__main__':
     unittest.main()