Added utility function to load all available key types. Updated docs. Updated import order so as not to break gevent/paramiko integration

Author: Dan

pssh/__init__.py

@@ -27,9 +27,10 @@
 
 See :mod:`pssh.ParallelSSHClient` and :mod:`pssh.SSHClient` for class documentation.
 """
-from .utils import enable_host_logger
+
 from .pssh_client import ParallelSSHClient
 from .ssh_client import SSHClient
+from .utils import enable_host_logger
 from .exceptions import UnknownHostException, \
      AuthenticationException, ConnectionErrorException, SSHException
 import logging

pssh/pssh_client.py

@@ -50,7 +50,7 @@ def __init__(self, hosts,
                  user=None, password=None, port=None, pkey=None,
                  forward_ssh_agent=True, num_retries=DEFAULT_RETRIES, timeout=120,
                  pool_size=10, proxy_host=None, proxy_port=22,
-                 agent=None, host_config=None):
+                 agent=None, host_config=None, channel_timeout=None):
         """
         :param hosts: Hosts to connect to
         :type hosts: list(str)
@@ -76,10 +76,9 @@ def __init__(self, hosts,
         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 number of hosts or 10, \
-        whichever is lower. Pool size will be *equal to* number of hosts if number\
-        of hosts is lower than the pool size specified as that would only \
-        increase overhead with no benefits.
+        hosts to execute tasks in parallel. Defaults to 10. Values over 500 \
+        are not likely to increase performance due to overhead in the single \
+        thread running our event loop.
         :type pool_size: int
         :param proxy_host: (Optional) SSH host to tunnel connection through \
         so that SSH clients connect to self.host via client -> proxy_host -> \
@@ -94,6 +93,9 @@ def __init__(self, hosts,
         :param host_config: (Optional) Per-host configuration for cases where \
         not all hosts use the same configuration values.
         :type host_config: dict
+        :param channel_timeout: (Optional) Time in seconds before an SSH operation \
+        times out.
+        :type channel_timeout: int
         
         **Example Usage**
         
@@ -244,6 +246,7 @@ def __init__(self, hosts,
         self.host_clients = {}
         self.agent = agent
         self.host_config = host_config if host_config else {}
+        self.channel_timeout = channel_timeout
 
     def run_command(self, *args, **kwargs):
         """Run command on all hosts in parallel, honoring self.pool_size,
@@ -448,7 +451,7 @@ def _get_host_config_values(self, host):
         _password = self.host_config.get(host, {}).get('password', self.password)
         _pkey = self.host_config.get(host, {}).get('private_key', self.pkey)
         return _user, _port, _password, _pkey
-
+    
     def _exec_command(self, host, *args, **kwargs):
         """Make SSHClient, run command on host"""
         if not host in self.host_clients or not self.host_clients[host]:
@@ -461,7 +464,8 @@ def _exec_command(self, host, *args, **kwargs):
                                                 timeout=self.timeout,
                                                 proxy_host=self.proxy_host,
                                                 proxy_port=self.proxy_port,
-                                                agent=self.agent)
+                                                agent=self.agent,
+                                                channel_timeout=self.channel_timeout)
         return self.host_clients[host].exec_command(*args, **kwargs)
 
     def get_output(self, cmd, output):
@@ -652,6 +656,7 @@ def _copy_file(self, host, local_file, remote_file, recurse=False):
                 timeout=self.timeout,
                 proxy_host=self.proxy_host,
                 proxy_port=self.proxy_port,
-                agent=self.agent)
+                agent=self.agent,
+                channel_timeout=self.channel_timeout)
         return self.host_clients[host].copy_file(local_file, remote_file,
                                                  recurse=recurse)

pssh/ssh_client.py

@@ -20,15 +20,14 @@
 
 import sys
 from gevent import sleep
-import logging
 import paramiko
 from paramiko.ssh_exception import ChannelException as channel_exception
 import os
 from socket import gaierror as sock_gaierror, error as sock_error
 from .exceptions import UnknownHostException, AuthenticationException, \
      ConnectionErrorException, SSHException
 from .constants import DEFAULT_RETRIES
-
+import logging
 
 host_logger = logging.getLogger('pssh.host_logger')
 logger = logging.getLogger(__name__)
@@ -43,7 +42,7 @@ def __init__(self, host,
                  user=None, password=None, port=None,
                  pkey=None, forward_ssh_agent=True,
                  num_retries=DEFAULT_RETRIES, agent=None, timeout=10,
-                 proxy_host=None, proxy_port=22):
+                 proxy_host=None, proxy_port=22, channel_timeout=None):
         """Connect to host honouring any user set configuration in ~/.ssh/config \
         or /etc/ssh/ssh_config
         
@@ -81,6 +80,9 @@ def __init__(self, host,
         :param proxy_port: (Optional) SSH port to use to login to proxy host if \
         set. Defaults to 22.
         :type proxy_port: int
+        :param channel_timeout: (Optional) Time in seconds before an SSH operation \
+        times out.
+        :type channel_timeout: int
         """
         ssh_config = paramiko.SSHConfig()
         _ssh_config_file = os.path.sep.join([os.path.expanduser('~'),
@@ -109,6 +111,7 @@ def __init__(self, host,
             self.client._agent = agent
         self.num_retries = num_retries
         self.timeout = timeout
+        self.channel_timeout = channel_timeout
         self.proxy_host, self.proxy_port = proxy_host, proxy_port
         self.proxy_client = None
         if self.proxy_host and self.proxy_port:
@@ -117,7 +120,7 @@ def __init__(self, host,
             self._connect_tunnel()
         else:
             self._connect(self.client, self.host, self.port)
-
+    
     def _connect_tunnel(self):
         """Connects to SSH server via an intermediate SSH tunnel server.
         client (me) -> tunnel (ssh server to proxy through) -> \
@@ -211,8 +214,8 @@ def exec_command(self, command, sudo=False, user=None,
         if self.forward_ssh_agent:
             agent_handler = paramiko.agent.AgentRequestHandler(channel)
         channel.get_pty()
-        # if self.timeout:
-        #     channel.settimeout(self.timeout)
+        if self.channel_timeout:
+            channel.settimeout(self.channel_timeout)
         _stdout, _stderr = channel.makefile('rb'), \
                            channel.makefile_stderr('rb')
         stdout, stderr = self._read_output_buffer(_stdout,), \

pssh/utils.py

@@ -20,6 +20,11 @@
 
 
 import logging
+import gevent
+from paramiko.rsakey import RSAKey
+from paramiko.dsskey import DSSKey
+from paramiko.ecdsakey import ECDSAKey
+from paramiko import SSHException
 
 host_logger = logging.getLogger('pssh.host_logger')
 logger = logging.getLogger('pssh')
@@ -42,6 +47,23 @@ def enable_host_logger():
     as it becomes available"""
     enable_logger(host_logger)
 
+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)
+    for keytype in [RSAKey, DSSKey, ECDSAKey]:
+        try:
+            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
+
 # def enable_pssh_logger():
 #     """Enable parallel-ssh's logger to stdout"""
 #     enable_logger(logger)