Added providing client private key as a PKey class to SSH clients. Added fail auth flag to fake server password authentication. Made fake server use gevent. Changed tests to use pre-generated test user private key

Author: pkittenis

fake_server/fake_server.py

@@ -8,15 +8,12 @@
 paramiko repository
 """
 
-# import multiprocessing
-import threading
-# import gevent
-# from gevent import monkey
-# monkey.patch_all()
+import gevent
+from gevent import monkey
+monkey.patch_all()
 import os
-import socket
-# from gevent import socket
-from threading import Event
+from gevent import socket
+from gevent.event import Event
 import sys
 import traceback
 import logging
@@ -38,6 +35,7 @@ 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
         return paramiko.AUTH_SUCCESSFUL
 
     def check_auth_publickey(self, username, key):
@@ -152,11 +150,7 @@ def handle_ssh_connection(cmd_req_response, sock, fail_auth = False):
         return
 
 def start_server(cmd_req_response, sock, fail_auth=False):
-    t = threading.Thread(target=listen, args=(cmd_req_response, sock,),
-                         kwargs={'fail_auth' : fail_auth})
-    t.start()
-    return t
-    # return gevent.spawn(listen, cmd_req_response, sock, fail_auth=fail_auth)
+    return gevent.spawn(listen, cmd_req_response, sock, fail_auth=fail_auth)
 
 if __name__ == "__main__":
     logging.basicConfig()

pssh.py

@@ -52,7 +52,8 @@ class SSHClient(object):
     overrides"""
 
     def __init__(self, host,
-                 user=None, password=None, port=None):
+                 user=None, password=None, port=None,
+                 pkey=None):
         """Connect to host honouring any user set configuration in ~/.ssh/config \
         or /etc/ssh/ssh_config
          
@@ -87,6 +88,7 @@ def __init__(self, host,
         self.channel = None
         self.user = user
         self.password = password
+        self.pkey = pkey
         self.port = port if port else 22
         self.host = resolved_address
         self._connect()
@@ -95,7 +97,8 @@ def _connect(self, retries=1):
         """Connect to host, throw UnknownHost exception on DNS errors"""
         try:
             self.client.connect(self.host, username=self.user,
-                                password=self.password, port=self.port)
+                                password=self.password, port=self.port,
+                                pkey=self.pkey)
         except socket.gaierror, e:
             logger.error("Could not resolve host '%s'", self.host,)
             while retries < NUM_RETRIES:
@@ -201,12 +204,11 @@ def copy_file(self, local_file, remote_file):
                         local_file, self.host, remote_file)
 
 class ParallelSSHClient(object):
-    """
-    Uses :mod:`pssh.SSHClient`, performs tasks over SSH on multiple hosts in \
+    """Uses :mod:`pssh.SSHClient`, performs tasks over SSH on multiple hosts in \
     parallel"""
 
     def __init__(self, hosts,
-                 user=None, password=None, port=None,
+                 user=None, password=None, port=None, pkey=None,
                  pool_size=10):
         """
         :param hosts: Hosts to connect to
@@ -265,6 +267,7 @@ def __init__(self, hosts,
         self.user = user
         self.password = password
         self.port = port
+        self.pkey = pkey
         # To hold host clients
         self.host_clients = dict((host, None) for host in hosts)
 
@@ -304,7 +307,7 @@ def _exec_command(self, host, *args, **kwargs):
         if not self.host_clients[host]:
             self.host_clients[host] = SSHClient(host, user=self.user,
                                                 password=self.password,
-                                                port=self.port)
+                                                port=self.port, pkey=self.pkey)
         return self.host_clients[host].exec_command(*args, **kwargs)
 
     def get_stdout(self, greenlet):

tests/test_pssh_client.py

@@ -5,25 +5,37 @@
 import unittest
 from pssh import ParallelSSHClient, UnknownHostException, \
     AuthenticationException, ConnectionErrorException, _setup_logger
-from fake_server.fake_server import start_server, make_socket, logger as server_logger
+from fake_server.fake_server import start_server, make_socket, logger as server_logger, \
+    paramiko_logger
 import random
 import logging
 import gevent
 import threading
+import paramiko
+import os
 
-_setup_logger(server_logger)
+# _setup_logger(server_logger)
+# _setup_logger(paramiko_logger)
+
+USER_KEY = paramiko.RSAKey.from_private_key_file(
+    os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key']))
 
 class ParallelSSHClientTest(unittest.TestCase):
 
     def setUp(self):
         self.fake_cmd = 'fake cmd'
         self.fake_resp = 'fake response'
+        self.user_key = USER_KEY
+        self.listen_socket = make_socket('127.0.0.1')
+        self.listen_port = self.listen_socket.getsockname()[1]
 
+    def tearDown(self):
+        del self.listen_socket
+        
     def test_pssh_client_exec_command(self):
-        sock = make_socket('127.0.0.1')
-        listen_port = sock.getsockname()[1]
-        server = start_server({ self.fake_cmd : self.fake_resp }, sock)
-        client = ParallelSSHClient(['127.0.0.1'], port=listen_port)
+        server = start_server({ self.fake_cmd : self.fake_resp }, self.listen_socket)
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   pkey=self.user_key)
         cmd = client.exec_command(self.fake_cmd)[0]
         output = client.get_stdout(cmd)
         expected = {'127.0.0.1' : {'exit_code' : 0}}
@@ -33,11 +45,10 @@ def test_pssh_client_exec_command(self):
         server.join()
 
     def test_pssh_client_auth_failure(self):
-        sock = make_socket('127.0.0.1')
-        listen_port = sock.getsockname()[1]
         server = start_server({ self.fake_cmd : self.fake_resp },
-                              sock, fail_auth=True)
-        client = ParallelSSHClient(['127.0.0.1'], port=listen_port)
+                              self.listen_socket, fail_auth=True)
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   pkey=self.user_key)
         cmd = client.exec_command(self.fake_cmd)[0]
         # Handle exception
         try:
@@ -47,3 +58,17 @@ def test_pssh_client_auth_failure(self):
             pass
         del client
         server.join()
+
+    def test_pssh_client_exec_command_password(self):
+        """Test password authentication. Fake server accepts any password
+        even empty string"""
+        server = start_server({ self.fake_cmd : self.fake_resp }, self.listen_socket)
+        client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
+                                   password='')
+        cmd = client.exec_command(self.fake_cmd)[0]
+        output = client.get_stdout(cmd)
+        expected = {'127.0.0.1' : {'exit_code' : 0}}
+        self.assertEqual(expected, output,
+                         msg = "Got unexpected command output - %s" % (output,))
+        del client
+        server.join()