Suppress gevent printing exceptions to stderr - we are handling them in ParallelSSH code. Resolves #30

Author: Dan

fake_server/fake_server.py

@@ -47,21 +47,29 @@
 host_key = paramiko.RSAKey(filename = os.path.sep.join([os.path.dirname(__file__), 'rsa.key']))
 
 class Server (paramiko.ServerInterface):
-    def __init__(self, transport, cmd_req_response = {}, fail_auth=False):
+    def __init__(self, transport, cmd_req_response = {}, fail_auth=False,
+                 ssh_exception=False):
         self.event = Event()
         self.cmd_req_response = cmd_req_response
         self.fail_auth = fail_auth
+        self.ssh_exception = ssh_exception
         self.transport = transport
 
     def check_channel_request(self, kind, chanid):
         return paramiko.OPEN_SUCCEEDED
 
     def check_auth_password(self, username, password):
-        if self.fail_auth: return paramiko.AUTH_FAILED
+        if self.fail_auth:
+            return paramiko.AUTH_FAILED
+        if self.ssh_exception:
+            raise paramiko.SSHException()
         return paramiko.AUTH_SUCCESSFUL
 
     def check_auth_publickey(self, username, key):
-        if self.fail_auth: return paramiko.AUTH_FAILED
+        if self.fail_auth:
+            return paramiko.AUTH_FAILED
+        if self.ssh_exception:
+            raise paramiko.SSHException()
         return paramiko.AUTH_SUCCESSFUL
 
     def get_allowed_auths(self, username):
@@ -131,7 +139,7 @@ def make_socket(listen_ip, port=0):
         return
     return sock
 
-def listen(cmd_req_response, sock, fail_auth=False,
+def listen(cmd_req_response, sock, fail_auth=False, ssh_exception=False,
            timeout=None):
     """Run a fake ssh server and given a cmd_to_run, send given \
     response to client connection. Returns (server, socket) tuple \
@@ -149,17 +157,18 @@ def listen(cmd_req_response, sock, fail_auth=False,
         traceback.print_exc()
         return
     handle_ssh_connection(cmd_req_response, sock, fail_auth=fail_auth,
-                          timeout=timeout)
+                          timeout=timeout, ssh_exception=ssh_exception)
 
-def _handle_ssh_connection(cmd_req_response, transport, fail_auth=False):
+def _handle_ssh_connection(cmd_req_response, transport, fail_auth=False,
+                           ssh_exception=False):
     try:
         transport.load_server_moduli()
     except:
         return
     transport.add_server_key(host_key)
     transport.set_subsystem_handler('sftp', paramiko.SFTPServer, StubSFTPServer)
     server = Server(transport, cmd_req_response=cmd_req_response,
-                    fail_auth=fail_auth)
+                    fail_auth=fail_auth, ssh_exception=ssh_exception)
     try:
         transport.start_server(server=server)
     except paramiko.SSHException, e:
@@ -180,7 +189,7 @@ def _handle_ssh_connection(cmd_req_response, transport, fail_auth=False):
     channel.close()
 
 def handle_ssh_connection(cmd_req_response, sock,
-                          fail_auth=False,
+                          fail_auth=False, ssh_exception=False,
                           timeout=None):
     conn, addr = sock.accept()
     logger.info('Got connection..')
@@ -190,7 +199,8 @@ def handle_ssh_connection(cmd_req_response, sock,
         gevent.Timeout(timeout).start()
     try:
         transport = paramiko.Transport(conn)
-        _handle_ssh_connection(cmd_req_response, transport, fail_auth=fail_auth)
+        _handle_ssh_connection(cmd_req_response, transport, fail_auth=fail_auth,
+                               ssh_exception=ssh_exception)
     except Exception, e:
         logger.error('*** Caught exception: %s: %s' % (str(e.__class__), str(e),))
         traceback.print_exc()
@@ -200,10 +210,10 @@ def handle_ssh_connection(cmd_req_response, sock,
             pass
         return
 
-def start_server(cmd_req_response, sock, fail_auth=False,
+def start_server(cmd_req_response, sock, fail_auth=False, ssh_exception=False,
                  timeout=None):
     return gevent.spawn(listen, cmd_req_response, sock, fail_auth=fail_auth,
-                        timeout=timeout)
+                        timeout=timeout, ssh_exception=ssh_exception)
 
 if __name__ == "__main__":
     logging.basicConfig()

pssh.py

@@ -31,6 +31,8 @@
 from gevent import monkey
 monkey.patch_all()
 import gevent.pool
+import gevent.hub
+gevent.hub.Hub.NOT_ERROR=(Exception,)
 import warnings
 from socket import gaierror as sock_gaierror, error as sock_error
 import logging
@@ -202,12 +204,12 @@ def _connect(self, client, host, port, sock=None, retries=1):
                                            str(error_type), self.host, self.port,
                                            retries, self.num_retries,)
         except paramiko.AuthenticationException, ex:
-            raise AuthenticationException(ex)
+            raise AuthenticationException(ex.message)
         # SSHException is more general so should be below other types
         # of SSH failure
         except paramiko.SSHException, ex:
             logger.error("General SSH error - %s", ex)
-            raise SSHException(ex)
+            raise SSHException(ex.message)
 
     def exec_command(self, command, sudo=False, user=None, **kwargs):
         """Wrapper to :mod:`paramiko.SSHClient.exec_command`

tests/test_pssh_client.py

@@ -21,7 +21,7 @@
 
 import unittest
 from pssh import ParallelSSHClient, UnknownHostException, \
-     AuthenticationException, ConnectionErrorException, logger as pssh_logger
+     AuthenticationException, ConnectionErrorException, SSHException, logger as pssh_logger
 from fake_server.fake_server import start_server, make_socket, \
      logger as server_logger, paramiko_logger
 import random
@@ -190,6 +190,21 @@ def test_pssh_client_auth_failure(self):
         del client
         server.join()
 
+    def test_pssh_client_ssh_exception(self):
+        server = start_server({ self.fake_cmd : self.fake_resp },
+                              self.listen_socket,
+                              ssh_exception=True)
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   pkey=self.user_key)
+        # Handle exception
+        try:
+            client.run_command(self.fake_cmd)
+            raise Exception("Expected SSHException, got none")
+        except SSHException, ex:
+            pass
+        del client
+        server.join()
+
     def test_pssh_client_timeout(self):
         server_timeout=0.2
         client_timeout=server_timeout-0.1