Added unittests for loading all possible key types plus invalid key. Updated load private key function. Added unittest for per host configuration in pssh client. Resolves #20

Author: Dan

pssh/pssh_client.py

@@ -201,11 +201,21 @@ def __init__(self, hosts,
         
         **Per-Host configuration**
         
+        Per host configuration can be provided for any of user, password port
+        and private key. Private key value is a :mod:`paramiko.PKey` object as
+        returned by :mod:`pssh.utils.load_private_key`.
         
+        :mod:`pssh.utils.load_private_key` accepts both file names and file-like
+        objects and will attempt to load all available key types, returning
+        `None` if they all fail.
+        
+        >>> from pssh.utils import load_private_key
         >>> host_config = { 'host1' : {'user': 'user1', 'password': 'pass',
-        ...                            'port': 2222, 'private_key': 'my_key.pem'},
+        ...                            'port': 2222,
+        ...                            'private_key': load_private_key('my_key.pem')},
         ...                 'host2' : {'user': 'user2', 'password': 'pass',
-        ...                            'port': 2223, 'private_key': 'my_other_key.pem'},
+        ...                            'port': 2223,
+        ...                            'private_key': load_private_key(open('my_other_key.pem'))},
         ...                 }
         >>> hosts = host_config.keys()
         >>> client = ParallelSSHClient(hosts, host_config=host_config)

pssh/utils.py

@@ -47,22 +47,22 @@ def enable_host_logger():
     as it becomes available"""
     enable_logger(host_logger)
 
-def load_private_key(pkey):
+def load_private_key(_pkey):
     """Load private key from pkey file object or filename
     
     :param pkey: File object or file name containing private key
     :type pkey: file/str"""
-    pkey = None
-    if not isinstance(pkey, file):
-        pkey = open(pkey)
+    if not hasattr(_pkey, 'read'):
+        _pkey = open(_pkey)
     for keytype in [RSAKey, DSSKey, ECDSAKey]:
         try:
-            pkey = keytype.from_private_key(pkey)
+            pkey = keytype.from_private_key(_pkey)
         except SSHException:
-            pass
-    if not pkey:
-        logger.error("Failed to load private key using all available key types - giving up..")
-    return pkey
+            _pkey.seek(0)
+            continue
+        else:
+            return pkey
+    logger.error("Failed to load private key using all available key types - giving up..")
 
 # def enable_pssh_logger():
 #     """Enable parallel-ssh's logger to stdout"""

tests/test_client_private_key_dsa

@@ -0,0 +1,12 @@
+-----BEGIN DSA PRIVATE KEY-----
+MIIBvQIBAAKBgQDnUBhCcJYWWxhFkD8RgFLNGhW59FILG819HFPCqFKoOzJAHsSp
+NRqlNi+IVxrFitFAfs8W/CwpNCZ3Co5I9dQ1pfeWs1PSrP2r7v95kJJBXSWh3c79
+416p0lS+Al2Ii4uNOyH25p+8I4W6VEm4YD0vpQuY5B3o8z0gGlIMWF49qwIVAMo4
+XAaxdeiQS1sg5ySMKxY0LKkhAoGBAM9+7NvijLjw3IeJes20DWyHE4AZMLnRhT6v
+Eo8xsy4yv6dHHkm0uxhO8wvtY7vtjDLPGjGI8kbqY1npO9ZHq32BFbmKQW7Lett+
+X7EoL7U/F1VAbxv6BtRkGslf3E6JqyfaHoxhsS/PvX9dA7MK4Jev55/86o68ETO2
+1PAu1SumAoGBAIKKFRzjN0jhSaMZyY9tBcW+94RcAaQv/jfij0EVYo+5F30y84NM
+mH7JhYbQ5A3UiAHyQSnz7q7vkAZ8gXH9NjVSu8yyMIvsW71gK3IHTC6rFWu9l1e6
+XlMmqblX/KY6z+K2Xg3zQ3djqRBTlNycGFwkLoM8I6n1nUDqLo/GI2JrAhUApeOU
+f0nVYotZrqSiJ168QM1Spzk=
+-----END DSA PRIVATE KEY-----

tests/test_client_private_key_ecdsa

@@ -0,0 +1,5 @@
+-----BEGIN EC PRIVATE KEY-----
+MHcCAQEEIMeaGUuH+Z+9BFhohwC+n08uzZMOOrdsdfUj5p4MIJGjoAoGCCqGSM49
+AwEHoUQDQgAE6x0VXeKfPMWkNfez96DZZNAPd7mYaetkl14rd9U4D4vqmwKCaQoZ
+Sar9yhqd2gS7/GiwVsxv7ICZnuZPBtVYJA==
+-----END EC PRIVATE KEY-----

tests/test_pssh_client.py

@@ -22,6 +22,7 @@
 import unittest
 from pssh import ParallelSSHClient, UnknownHostException, \
      AuthenticationException, ConnectionErrorException, SSHException, logger as pssh_logger
+from pssh.utils import load_private_key
 from embedded_server.embedded_server import start_server, make_socket, \
      logger as server_logger, paramiko_logger
 from embedded_server.fake_agent import FakeAgent
@@ -34,8 +35,8 @@
 import warnings
 import shutil
 
-USER_KEY = paramiko.RSAKey.from_private_key_file(
-    os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key']))
+PKEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key'])
+USER_KEY = paramiko.RSAKey.from_private_key_file(PKEY_FILENAME)
 
 server_logger.setLevel(logging.DEBUG)
 pssh_logger.setLevel(logging.DEBUG)
@@ -648,3 +649,44 @@ def test_escaped_quotes(self):
         self.assertEqual(expected, stdout,
                          msg="Got unexpected output. Expected %s, got %s" % (
                              expected, stdout,))
+
+    def test_host_config(self):
+        """Test per-host configuration functionality of ParallelSSHClient"""
+        hosts = ['127.0.0.%01d' % n for n in xrange(1,3)]
+        host_config = dict.fromkeys(hosts)
+        servers = []
+        user = 'overriden_user'
+        password = 'overriden_pass'
+        for host in hosts:
+            _socket = make_socket(host)
+            port = _socket.getsockname()[1]
+            host_config[host] = {}
+            host_config[host]['port'] = port
+            host_config[host]['user'] = user
+            host_config[host]['password'] = password
+            server = start_server(_socket, fail_auth=hosts.index(host))
+            servers.append((server, port))
+        pkey_data = load_private_key(PKEY_FILENAME)
+        host_config[hosts[0]]['private_key'] = pkey_data
+        client = ParallelSSHClient(hosts, host_config=host_config)
+        output = client.run_command(self.fake_cmd, stop_on_errors=False)
+        client.join(output)
+        for host in hosts:
+            self.assertTrue(host in output)
+        try:
+            raise output[hosts[1]]['exception']
+        except AuthenticationException, ex:
+            pass
+        else:
+            raise AssertionError("Expected AutnenticationException on host %s",
+                                 hosts[0])
+        self.assertFalse(output[hosts[1]]['exit_code'],
+                         msg="Execution failed on host %s" % (hosts[1],))
+        self.assertTrue(client.host_clients[hosts[0]].user == user,
+                        msg="Host config user override failed")
+        self.assertTrue(client.host_clients[hosts[0]].password == password,
+                        msg="Host config password override failed")
+        self.assertTrue(client.host_clients[hosts[0]].pkey == pkey_data,
+                        msg="Host config pkey override failed")
+        for (server, _) in servers:
+            server.kill()

tests/test_utils.py

@@ -1,6 +1,11 @@
 from pssh import utils
 import unittest
+import os
+from cStringIO import StringIO
 
+PKEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key'])
+DSA_KEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key_dsa'])
+ECDSA_KEY_FILENAME = os.path.sep.join([os.path.dirname(__file__), 'test_client_private_key_ecdsa'])
 
 class ParallelSSHUtilsTest(unittest.TestCase):
 
@@ -13,3 +18,13 @@ def test_enabling_host_logger(self):
     def test_enabling_pssh_logger(self):
         utils.enable_logger(utils.logger)
         self.assertTrue(len(utils.logger.handlers)==1)
+
+
+    def test_loading_key_files(self):
+        for key_filename in [PKEY_FILENAME, DSA_KEY_FILENAME, ECDSA_KEY_FILENAME]:
+            pkey = utils.load_private_key(key_filename)
+            self.assertTrue(pkey, msg="Error loading key from file %s" % (key_filename,))
+            pkey = utils.load_private_key(open(key_filename))
+            self.assertTrue(pkey, msg="Error loading key from open file object for file %s" % (key_filename,))
+        fake_key = StringIO("blah blah fakey fakey key")
+        self.assertFalse(utils.load_private_key(fake_key))