Added keepalive functionality and integration tests.

Updated changelog.

Author: Pan

Changelog.rst

@@ -9,6 +9,7 @@ Changes
 
 * Removed libssh2 native library dependency in favour of bundled ssh2-python libssh2 library.
 * Changed native client forward agent default behaviour to off due to incompatibility with certain SSH server implementations.
+* Added keep-alive functionality to native client - defaults to ``60`` seconds. ``ParallelSSHClient.keepalive_seconds`` to configure interval, ``0`` to disable.
 
 
 1.8.2

pssh/clients/native/parallel.py

@@ -39,7 +39,8 @@ def __init__(self, hosts, user=None, password=None, port=22, pkey=None,
                  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,
-                 forward_ssh_agent=False, tunnel_timeout=None):
+                 forward_ssh_agent=False, tunnel_timeout=None,
+                 keepalive_seconds=60):
         """
         :param hosts: Hosts to connect to
         :type hosts: list(str)
@@ -122,6 +123,7 @@ def __init__(self, hosts, user=None, password=None, port=22, pkey=None,
         self._tunnel_lock = None
         self._tunnel_timeout = tunnel_timeout
         self._clients_lock = RLock()
+        self.keepalive_seconds = keepalive_seconds
 
     def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
                     use_pty=False, host_args=None, shell=None,
@@ -381,7 +383,8 @@ def _make_ssh_client(self, host):
                     timeout=self.timeout,
                     allow_agent=self.allow_agent, retry_delay=self.retry_delay,
                     proxy_host=proxy_host, _auth_thread_pool=auth_thread_pool,
-                    forward_ssh_agent=self.forward_ssh_agent)
+                    forward_ssh_agent=self.forward_ssh_agent,
+                    keepalive_seconds=self.keepalive_seconds)
 
     def copy_file(self, local_file, remote_file, recurse=False, copy_args=None):
         """Copy local file to remote file in parallel

pssh/clients/native/single.py

@@ -26,7 +26,7 @@
 from socket import gaierror as sock_gaierror, error as sock_error
 from warnings import warn
 
-from gevent import sleep, socket, get_hub
+from gevent import sleep, socket, get_hub, spawn
 from gevent.hub import Hub
 from ssh2.error_codes import LIBSSH2_ERROR_EAGAIN
 from ssh2.exceptions import SFTPHandleError, SFTPProtocolError, \
@@ -68,7 +68,7 @@ def __init__(self, host,
                  allow_agent=True, timeout=None,
                  forward_ssh_agent=False,
                  proxy_host=None,
-                 _auth_thread_pool=True):
+                 _auth_thread_pool=True, keepalive_seconds=60):
         """:param host: Host name or IP to connect to.
         :type host: str
         :param user: User to connect as. Defaults to logged in user.
@@ -100,6 +100,8 @@ def __init__(self, host,
         :param proxy_host: Connection to host is via provided proxy host
           and client should use self.proxy_host for connection attempts.
         :type proxy_host: str
+        :param keepalive_seconds: Interval of keep alive messages being sent to
+          server. Set to ``0`` or ``False`` to disable.
 
         :raises: :py:class:`pssh.exceptions.PKeyFileError` on errors finding
           provided private key.
@@ -123,6 +125,8 @@ def __init__(self, host,
         self._host = proxy_host if proxy_host else host
         self.pkey = _validate_pkey_path(pkey, self.host)
         self._connect(self._host, self.port)
+        self.keepalive_seconds = keepalive_seconds
+        self._keepalive_greenlet = None
         if _auth_thread_pool:
             THREAD_POOL.apply(self._init)
         else:
@@ -147,6 +151,18 @@ def __enter__(self):
     def __exit__(self, *args):
         self.disconnect()
 
+    def spawn_send_keepalive(self):
+        """Spawns a new greenlet that sends keep alive messages every
+        self.keepalive_seconds"""
+        return spawn(self._send_keepalive)
+
+    def _send_keepalive(self):
+        while True:
+            sleep(self._eagain(self.session.keepalive_send))
+
+    def configure_keepalive(self):
+        self.session.keepalive_config(False, self.keepalive_seconds)
+
     def _connect_init_retry(self, retries):
         retries += 1
         self.session = None
@@ -179,6 +195,9 @@ def _init(self, retries=1):
             msg = "Authentication error while connecting to %s:%s - %s"
             raise AuthenticationException(msg, self.host, self.port, ex)
         self.session.set_blocking(0)
+        if self.keepalive_seconds:
+            self.configure_keepalive()
+            self._keepalive_greenlet = self.spawn_send_keepalive()
 
     def _connect(self, host, port, retries=1):
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

tests/test_native_parallel_client.py

@@ -1401,3 +1401,11 @@ def test_disable_agent_forward(self):
             num_retries=1)
         client.join(client.run_command(self.cmd))
         self.assertFalse(client.host_clients[self.host].forward_ssh_agent)
+
+    def test_keepalive_off(self):
+        client = ParallelSSHClient(
+            [self.host], port=self.port, pkey=self.user_key,
+            keepalive_seconds=0,
+            num_retries=1)
+        client.join(client.run_command(self.cmd))
+        self.assertFalse(client.host_clients[self.host].keepalive_seconds)

tests/test_native_single_client.py

@@ -155,4 +155,4 @@ def test_connection_timeout(self):
                     num_retries=1, timeout=1, _auth_thread_pool=False)
         # Should fail within greenlet timeout, otherwise greenlet will
         # raise timeout which will fail the test
-        self.assertRaises(ConnectionErrorException, cmd.get, timeout=1.1)
+        self.assertRaises(ConnectionErrorException, cmd.get, timeout=2)