Added native clients proxying/tunnelling implementation, tests, exception.

Author: Pan

Changelog.rst

@@ -1,6 +1,17 @@
 Change Log
 ============
 
+1.3.0
+++++++
+
+Changes
+---------
+
+* Native clients proxy implementation
+* Native clients connection and authentication retry mechanism
+
+Proxy/tunnelling implementation is experimental - please report any issues.
+
 1.2.1
 ++++++
 

pssh/exceptions.py

@@ -57,3 +57,7 @@ class SFTPError(Exception):
 class SFTPIOError(SFTPError):
     """Raised on SFTP IO errors"""
     pass
+
+
+class ProxyError(Exception):
+    """Raised on proxy errors"""

pssh/pssh2_client.py

@@ -16,10 +16,13 @@
 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 
 import logging
+from gevent import sleep
 
 from .base_pssh import BaseParallelSSHClient
 from .constants import DEFAULT_RETRIES, RETRY_DELAY
 from .ssh2_client import SSHClient
+from .exceptions import ProxyError
+from .tunnel import Tunnel
 
 
 logger = logging.getLogger(__name__)
@@ -30,7 +33,9 @@ class ParallelSSHClient(BaseParallelSSHClient):
 
     def __init__(self, hosts, user=None, password=None, port=None, pkey=None,
                  num_retries=DEFAULT_RETRIES, timeout=None, pool_size=10,
-                 allow_agent=True, host_config=None, retry_delay=RETRY_DELAY):
+                 allow_agent=True, host_config=None, retry_delay=RETRY_DELAY,
+                 proxy_host=None, proxy_port=22,
+                 proxy_user=None, proxy_password=None, proxy_pkey=None):
         """
         :param hosts: Hosts to connect to
         :type hosts: list(str)
@@ -64,23 +69,47 @@ def __init__(self, hosts, user=None, password=None, port=None, pkey=None,
           not all hosts use the same configuration.
         :type host_config: dict
         :param allow_agent: (Optional) set to False to disable connecting to
-          the system's SSH agent
+          the system's SSH agent.
         :type allow_agent: bool
+        :param proxy_host: (Optional) SSH host to tunnel connection through
+          so that SSH clients connect to host via client -> proxy_host -> host
+        :type proxy_host: str
+        :param proxy_port: (Optional) SSH port to use to login to proxy host if
+          set. Defaults to 22.
+        :type proxy_port: int
+        :param proxy_user: (Optional) User to login to ``proxy_host`` as.
+          Defaults to logged in user.
+        :type proxy_user: str
+        :param proxy_password: (Optional) Password to login to ``proxy_host``
+          with. Defaults to no password.
+        :type proxy_password: str
+        :param proxy_pkey: (Optional) Private key file to be used for
+          authentication with ``proxy_host``. Defaults to available keys from
+          SSHAgent and user's SSH identities.
+        :type proxy_pkey: Private key file path to use. Note that the public
+          key file pair *must* also exist in the same location with name
+          ``<pkey>.pub``.
         """
         BaseParallelSSHClient.__init__(
             self, hosts, user=user, password=password, port=port, pkey=pkey,
             allow_agent=allow_agent, num_retries=num_retries,
             timeout=timeout, pool_size=pool_size,
             host_config=host_config, retry_delay=retry_delay)
+        self.proxy_host = proxy_host
+        self.proxy_port = proxy_port
+        self.proxy_pkey = proxy_pkey
+        self.proxy_user = proxy_user
+        self.proxy_password = proxy_password
+        self._tunnels = {}
 
     def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
                     use_pty=False, host_args=None, shell=None,
                     encoding='utf-8'):
         """Run command on all hosts in parallel, honoring self.pool_size,
         and return output dictionary.
 
-        This function will block until all commands have been successfully
-        received by remote servers and then return immediately.
+        This function will block until all commands have been received
+        by remote servers and then return immediately.
 
         More explicitly, function will return after connection and
         authentication establishment and after commands have been accepted by
@@ -139,7 +168,10 @@ def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
         :raises: :py:class:`TypeError` on not enough host arguments for cmd
           string format
         :raises: :py:class:`KeyError` on no host argument key in arguments
-          dict for cmd string format"""
+          dict for cmd string format
+        :raises: :py:class:`pssh.exceptions.ProxyErrors` on errors connecting
+          to proxy if a proxy host has been set.
+        """
         return BaseParallelSSHClient.run_command(
             self, command, stop_on_errors=stop_on_errors, host_args=host_args,
             user=user, shell=shell, sudo=sudo,
@@ -184,11 +216,34 @@ def _get_exit_code(self, channel):
             return
         return channel.get_exit_status()
 
+    def _start_tunnel(self, host):
+        tunnel = Tunnel(
+            self.proxy_host, host, self.port, user=self.proxy_user,
+            password=self.proxy_password, port=self.proxy_port,
+            pkey=self.proxy_pkey, num_retries=self.num_retries,
+            timeout=self.timeout, retry_delay=self.retry_delay,
+            allow_agent=self.allow_agent)
+        tunnel.daemon = True
+        tunnel.start()
+        self._tunnels[host] = tunnel
+        while not tunnel.tunnel_open.is_set():
+            logger.debug("Waiting for tunnel to become active")
+            sleep(.1)
+            if not tunnel.is_alive():
+                msg = "Proxy authentication failed"
+                logger.error(msg)
+                raise ProxyError(msg)
+        return tunnel
+
     def _make_ssh_client(self, host):
         if host not in self.host_clients or self.host_clients[host] is None:
+            if self.proxy_host is not None:
+                tunnel = self._start_tunnel(host)
             _user, _port, _password, _pkey = self._get_host_config_values(host)
+            _host = host if self.proxy_host is None else '127.0.0.1'
+            _port = _port if self.proxy_host is None else tunnel.listen_port
             self.host_clients[host] = SSHClient(
-                host, user=_user, password=_password, port=_port, pkey=_pkey,
+                _host, user=_user, password=_password, port=_port, pkey=_pkey,
                 num_retries=self.num_retries, timeout=self.timeout,
                 allow_agent=self.allow_agent, retry_delay=self.retry_delay)
 

pssh/pssh_client.py

@@ -128,11 +128,11 @@ def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
         """Run command on all hosts in parallel, honoring self.pool_size,
         and return output buffers.
 
-        This function will block until all commands have been *sent* to remote
-        servers and then return immediately
+        This function will block until all commands have been received
+        by remote servers and then return immediately.
 
         More explicitly, function will return after connection and
-        authentication establishment and after commands have been sent to
+        authentication establishment and after commands have been received by
         successfully established SSH channels.
 
         Any connection and/or authentication exceptions will be raised here

pssh/ssh2_client.py

@@ -25,8 +25,7 @@
     WIN_PLATFORM = False
 from socket import gaierror as sock_gaierror, error as sock_error
 
-from gevent import sleep, get_hub
-from gevent import socket
+from gevent import sleep, socket, get_hub
 from gevent.hub import Hub
 from ssh2.error_codes import LIBSSH2_ERROR_EAGAIN
 from ssh2.exceptions import AuthenticationError, AgentError, \
@@ -89,63 +88,77 @@ def __init__(self, host,
           the system's SSH agent
         :type allow_agent: bool
         """
-        # proxy_host=None, proxy_port=22, proxy_user=None,
-        # proxy_password=None, proxy_pkey=None,
         self.host = host
         self.user = user if user else None
-        if self.user is None and WIN_PLATFORM is False:
+        if self.user is None and not WIN_PLATFORM:
             self.user = pwd.getpwuid(os.geteuid()).pw_name
-        elif self.user is None and WIN_PLATFORM is True:
+        elif self.user is None and WIN_PLATFORM:
             raise ValueError("Must provide user parameter on Windows")
         self.password = password
         self.port = port if port else 22
         self.pkey = pkey
         self.num_retries = num_retries
-        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.sock = None
         self.timeout = timeout * 1000 if timeout else None
         self.retry_delay = retry_delay
         self.allow_agent = allow_agent
-        self._connect()
+        self.session = None
+        self._connect(self.host, self.port)
         THREAD_POOL.apply(self._init)
 
-    def _init(self):
+    def _connect_init_retry(self, retries):
+        retries += 1
+        self.session = None
+        if not self.sock.closed:
+            self.sock.close()
+        sleep(self.retry_delay)
+        self._connect(self.host, self.port, retries=retries)
+        return self._init(retries=retries)
+
+    def _init(self, retries=1):
         self.session = Session()
         if self.timeout:
             self.session.set_timeout(self.timeout)
         try:
             self.session.handshake(self.sock)
         except SessionHandshakeError as ex:
+            while retries < self.num_retries:
+                return self._connect_init_retry(retries)
             msg = "Error connecting to host %s:%s - %s"
             raise SessionError(msg, self.host, self.port, ex)
         try:
             self.auth()
         except AuthenticationError as ex:
+            while retries < self.num_retries:
+                return self._connect_init_retry(retries)
             msg = "Authentication error while connecting to %s:%s - %s"
             raise AuthenticationException(msg, self.host, self.port, ex)
         self.session.set_blocking(0)
 
-    def _connect(self, retries=1):
+    def _connect(self, host, port, retries=1):
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        logger.debug("Connecting to %s:%s", host, port)
         try:
-            self.sock.connect((self.host, self.port))
+            self.sock.connect((host, port))
         except sock_gaierror as ex:
             logger.error("Could not resolve host '%s' - retry %s/%s",
-                         self.host, retries, self.num_retries)
+                         host, retries, self.num_retries)
             while retries < self.num_retries:
                 sleep(self.retry_delay)
-                return self._connect(retries=retries+1)
+                return self._connect(host, port, retries=retries+1)
             raise UnknownHostException("Unknown host %s - %s - retry %s/%s",
-                                       self.host, str(ex.args[1]), retries,
+                                       host, str(ex.args[1]), retries,
                                        self.num_retries)
         except sock_error as ex:
             logger.error("Error connecting to host '%s:%s' - retry %s/%s",
-                         self.host, self.port, retries, self.num_retries)
+                         host, port, retries, self.num_retries)
             while retries < self.num_retries:
                 sleep(self.retry_delay)
-                return self._connect(retries=retries+1)
+                return self._connect(host, port, retries=retries+1)
             error_type = ex.args[1] if len(ex.args) > 1 else ex.args[0]
             raise ConnectionErrorException(
                 "Error connecting to host '%s:%s' - %s - retry %s/%s",
-                self.host, self.port, str(error_type), retries,
+                host, port, str(error_type), retries,
                 self.num_retries,)
 
     def _pkey_auth(self):