Merge "Improve failover of SSH connections"

4 files changed, 106 insertions, 72 deletions

diff --git a/fuelweb_test/helpers/checkers.py b/fuelweb_test/helpers/checkers.py
index 2780b15..186f2ef 100644
--- a/fuelweb_test/helpers/checkers.py
+++ b/fuelweb_test/helpers/checkers.py
@@ -1466,8 +1466,8 @@ def check_package_version(ip, package_name, expected_version, condition='ge'):
     :param condition: predicate can be on of eq, ne, lt, le, ge, gt
     :return None: or raise UnexpectedExitCode
     """
-    cmd = "dpkg -s {0} " \
-          "| awk -F': ' '/Version/ {{print \$2}}'".format(package_name)
+    cmd = ("dpkg -s {0} "
+           "| awk -F': ' '/Version/ {{print $2}}'".format(package_name))
     logger.debug(cmd)
     result = ssh_manager.execute_on_remote(
         ip,
diff --git a/fuelweb_test/helpers/ssh_manager.py b/fuelweb_test/helpers/ssh_manager.py
index 143467f..63612c1 100644
--- a/fuelweb_test/helpers/ssh_manager.py
+++ b/fuelweb_test/helpers/ssh_manager.py
@@ -19,9 +19,10 @@ import traceback
 from warnings import warn
 
 from devops.helpers.metaclasses import SingletonMeta
+from devops.helpers.ssh_client import SSHAuth
 from devops.helpers.ssh_client import SSHClient
 from paramiko import RSAKey
-from paramiko.ssh_exception import AuthenticationException
+from paramiko import SSHException
 import six
 
 from fuelweb_test import logger
@@ -67,8 +68,15 @@ class SSHManager(six.with_metaclass(SingletonMeta, object)):
         self.slave_login = slave_login
         self.__slave_password = slave_password
 
-    @staticmethod
-    def _connect(remote):
+    def _get_keys(self):
+        keys = []
+        admin_remote = self.get_remote(self.admin_ip)
+        key_string = '/root/.ssh/id_rsa'
+        with admin_remote.open(key_string) as f:
+            keys.append(RSAKey.from_private_key(f))
+        return keys
+
+    def connect(self, remote):
         """ Check if connection is stable and return this one
 
         :param remote:
@@ -80,99 +88,129 @@ class SSHManager(six.with_metaclass(SingletonMeta, object)):
                     seconds=5,
                     error_message="Socket timeout! Forcing reconnection"):
                 remote.check_call("cd ~")
-        except:
+        except Exception:
             logger.debug(traceback.format_exc())
-            logger.info('SSHManager: Check for current connection fails. '
-                        'Trying to reconnect')
-            remote.reconnect()
+            logger.debug('SSHManager: Check for current connection fails. '
+                         'Trying to reconnect')
+            remote = self.reconnect(remote)
         return remote
 
-    def _get_keys(self):
-        keys = []
-        admin_remote = self.get_remote(self.admin_ip)
-        key_string = '/root/.ssh/id_rsa'
-        with admin_remote.open(key_string) as f:
-            keys.append(RSAKey.from_private_key(f))
-        return keys
+    def reconnect(self, remote):
+        """ Reconnect to remote or update connection
 
-    def get_remote(self, ip, port=22):
-        """ Function returns remote SSH connection to node by ip address
+        :param remote:
+        :return:
+        """
+        ip = remote.hostname
+        port = remote.port
+        try:
+            remote.reconnect()
+        except SSHException:
+            self.update_connection(ip=ip, port=port)
+        return self.connections[(ip, port)]
+
+    def init_remote(self, ip, port=22, custom_creds=None):
+        """ Initialise connection to remote
 
         :param ip: IP of host
         :type ip: str
         :param port: port for SSH
         :type port: int
-        :rtype: SSHClient
+        :param custom_creds: custom creds
+        :type custom_creds: dict
         """
-        if (ip, port) not in self.connections:
-            logger.debug('SSH_MANAGER: Create new connection for '
-                         '{ip}:{port}'.format(ip=ip, port=port))
+        logger.debug('SSH_MANAGER: Create new connection for '
+                     '{ip}:{port}'.format(ip=ip, port=port))
 
-            keys = self._get_keys() if ip != self.admin_ip else []
-            if ip == self.admin_ip:
+        keys = self._get_keys() if ip != self.admin_ip else []
+        if ip == self.admin_ip:
+            ssh_client = SSHClient(
+                host=ip,
+                port=port,
+                auth=SSHAuth(
+                    username=self.admin_login,
+                    password=self.__admin_password,
+                    keys=keys)
+            )
+            ssh_client.sudo_mode = SSH_FUEL_CREDENTIALS['sudo']
+        elif custom_creds:
+            ssh_client = SSHClient(
+                host=ip,
+                port=port,
+                auth=SSHAuth(**custom_creds))
+        else:
+            try:
                 ssh_client = SSHClient(
                     host=ip,
                     port=port,
-                    username=self.admin_login,
-                    password=self.__admin_password,
-                    private_keys=keys
-                )
-                ssh_client.sudo_mode = SSH_FUEL_CREDENTIALS['sudo']
-            else:
-                try:
-                    ssh_client = SSHClient(
-                        host=ip,
-                        port=port,
+                    auth=SSHAuth(
                         username=self.slave_login,
                         password=self.__slave_password,
-                        private_keys=keys
-                    )
-                except AuthenticationException:
-                    ssh_client = SSHClient(
-                        host=ip,
-                        port=port,
+                        keys=keys)
+                )
+            except SSHException:
+                ssh_client = SSHClient(
+                    host=ip,
+                    port=port,
+                    auth=SSHAuth(
                         username=self.slave_fallback_login,
                         password=self.__slave_password,
-                        private_keys=keys
-                    )
-                ssh_client.sudo_mode = SSH_SLAVE_CREDENTIALS['sudo']
-            self.connections[(ip, port)] = ssh_client
-        logger.debug('SSH_MANAGER: Return existed connection for '
-                     '{ip}:{port}'.format(ip=ip, port=port))
+                        keys=keys)
+                )
+            ssh_client.sudo_mode = SSH_SLAVE_CREDENTIALS['sudo']
+
+        self.connections[(ip, port)] = ssh_client
+        logger.debug('SSH_MANAGER: New connection for '
+                     '{ip}:{port} is created'.format(ip=ip, port=port))
+
+    def get_remote(self, ip, port=22):
+        """ Function returns remote SSH connection to node by ip address
+
+        :param ip: IP of host
+        :type ip: str
+        :param port: port for SSH
+        :type port: int
+        :rtype: SSHClient
+        """
+        if (ip, port) in self.connections:
+            logger.debug('SSH_MANAGER: Return existed connection for '
+                         '{ip}:{port}'.format(ip=ip, port=port))
+        else:
+            self.init_remote(ip=ip, port=port)
         logger.debug('SSH_MANAGER: Connections {0}'.format(self.connections))
-        return self._connect(self.connections[(ip, port)])
+        return self.connect(self.connections[(ip, port)])
 
-    def update_connection(self, ip, login=None, password=None,
-                          keys=None, port=22):
+    def update_connection(self, ip, port=22, login=None, password=None,
+                          keys=None):
         """Update existed connection
 
         :param ip: host ip string
+        :param port: ssh port int
         :param login: login string
         :param password: password string
         :param keys: list of keys
-        :param port: ssh port int
         :return: None
         """
         if (ip, port) in self.connections:
-            logger.info('SSH_MANAGER: Close connection for {ip}:{port}'.format(
-                ip=ip, port=port))
-            self.connections[(ip, port)].clear()
-            logger.info('SSH_MANAGER: Create new connection for '
-                        '{ip}:{port}'.format(ip=ip, port=port))
-
-        self.connections[(ip, port)] = SSHClient(
-            host=ip,
-            port=port,
-            username=login,
-            password=password,
-            private_keys=keys if keys is not None else []
-        )
+            logger.debug('SSH_MANAGER: Close connection for {ip}:{port}'
+                         .format(ip=ip, port=port))
+            ssh_client = self.connections.pop((ip, port))
+            ssh_client.close()
+        if login and (password or keys):
+            custom_creds = {
+                'username': login,
+                'password': password,
+                'keys': keys
+            }
+        else:
+            custom_creds = None
+        self.init_remote(ip=ip, port=port, custom_creds=custom_creds)
 
     def clean_all_connections(self):
         for (ip, port), connection in self.connections.items():
             connection.clear()
-            logger.info('SSH_MANAGER: Close connection for {ip}:{port}'.format(
-                ip=ip, port=port))
+            logger.debug('SSH_MANAGER: Close connection for {ip}:{port}'
+                         .format(ip=ip, port=port))
 
     def execute(self, ip, cmd, port=22, sudo=None):
         remote = self.get_remote(ip=ip, port=port)
diff --git a/fuelweb_test/models/environment.py b/fuelweb_test/models/environment.py
index 08dd32f..3eeb915 100644
--- a/fuelweb_test/models/environment.py
+++ b/fuelweb_test/models/environment.py
@@ -377,6 +377,7 @@ class EnvironmentModel(six.with_metaclass(SingletonMeta, object)):
             )
             self.ssh_manager.update_connection(
                 ip=self.ssh_manager.admin_ip,
+                port=22,
                 login=new_login,
                 password=new_password
             )
diff --git a/fuelweb_test/tests/test_multipath_devices.py b/fuelweb_test/tests/test_multipath_devices.py
index 0d07aac..3fc927f 100644
--- a/fuelweb_test/tests/test_multipath_devices.py
+++ b/fuelweb_test/tests/test_multipath_devices.py
@@ -25,7 +25,6 @@ from fuelweb_test.settings import MULTIPATH
 from fuelweb_test.settings import MULTIPATH_TEMPLATE
 from fuelweb_test.settings import NEUTRON_SEGMENT
 from fuelweb_test.settings import SLAVE_MULTIPATH_DISKS_COUNT
-from fuelweb_test.settings import SSH_FUEL_CREDENTIALS
 from fuelweb_test.settings import REPLACE_DEFAULT_REPOS
 from fuelweb_test.settings import REPLACE_DEFAULT_REPOS_ONLY_ONCE
 from fuelweb_test.tests import base_test_case
@@ -53,9 +52,7 @@ class TestMultipath(base_test_case.TestBasic):
         """
         cmd = "multipath -l -v2"
 
-        ssh_manager.update_connection(ip, SSH_FUEL_CREDENTIALS['login'],
-                                      SSH_FUEL_CREDENTIALS['password'],
-                                      keys=ssh_manager._get_keys())
+        ssh_manager.update_connection(ip)
         ssh_manager.get_remote(ip)
         result = ssh_manager.execute_on_remote(
             ip=ip,
@@ -106,9 +103,7 @@ class TestMultipath(base_test_case.TestBasic):
         """
         cmd = "lsblk -lo NAME,TYPE,MOUNTPOINT | grep '/$' | grep -c lvm"
 
-        ssh_manager.update_connection(ip, SSH_FUEL_CREDENTIALS['login'],
-                                      SSH_FUEL_CREDENTIALS['password'],
-                                      keys=ssh_manager._get_keys())
+        ssh_manager.update_connection(ip)
         ssh_manager.get_remote(ip)
         result = ssh_manager.execute_on_remote(
             ip=ip,