Reuse SSH Client

Allow the installer to re-use the paramiko SSH Client instead of creating a new one every time, resulting in slightly lower install time.

Regression: Tested by installing a AIO-SX system, PASS

Story: 2005051
Task: 48502

Change-Id: I1fcb0df2a0a04a2983c7f7f32bf12f3351c0e8d9
Signed-off-by: Felipe Freire <felipe.freire@encora.com>

virtualbox/pybox/exceptions/__init__.py

@@ -0,0 +1 @@
+from .ssh_exception import InvalidSSHConnection

virtualbox/pybox/exceptions/ssh_exception.py

@@ -0,0 +1,2 @@
+class InvalidSSHConnection(Exception):
+    pass

virtualbox/pybox/install_vbox.py

@@ -34,11 +34,14 @@ from consts.networking import NICs, OAM, MGMT, Serial
 from consts.timeout import HostTimeout
 from consts import env
 
+from exceptions import InvalidSSHConnection
+
 from Parser import handle_args
 
 
 # Global vars
 V_BOX_OPTIONS = None
+SSH_CONNECTIONS = {}
 
 # Network
 OAM_CONFIG = [getattr(OAM, attr) for attr in dir(OAM) if not attr.startswith('__')]
@@ -972,6 +975,21 @@ def _connect_to_ssh(node='floating'):
         LOG.error("#### Failed SSH connection\nError: %s", repr(exc))
         raise
 
+    SSH_CONNECTIONS[node] = ssh
+    return ssh
+
+def ssh_handler(node='floating'):
+    """
+    Handles the SSH connection. Tries to retrieve a already existing connection.
+    If it doesn't exist, or isn't active, creates a new one.
+    """
+    try:
+        ssh = SSH_CONNECTIONS[node]
+        ssh_transport = ssh.get_transport()
+        if ssh_transport is None or not ssh_transport.is_active():
+            raise InvalidSSHConnection()
+    except (InvalidSSHConnection, KeyError):
+        ssh = _connect_to_ssh(node)
     return ssh
 
 
@@ -1000,19 +1018,21 @@ def connect_to_ssh(node='floating'):
 
             Returns: return code of decorated function
             """
-
-            try:
-                ssh = _connect_to_ssh(node)
+            ssh = ssh_handler(node)
             kwargs['ssh_client'] = ssh
             return func(*args, **kwargs)
-            finally:
-                if ssh:
-                    ssh.close()
 
         return connect_to_ssh_wrapper
 
     return connect_to_ssh_decorator
 
+def close_ssh_connections(ssh_clients: dict):
+    """
+    Closes the connection with all current created SSH Clients.
+    """
+    for _, client in ssh_clients.items():
+        if client is not None:
+            client.close()
 
 def stage_test_success():
     """Prints a log message indicating the execution of a test stage."""
@@ -1790,7 +1810,7 @@ def run_custom_script(script, timeout, console, mode):
     LOG.info("     console mode:   %s", console)
     LOG.info("     user mode: %s", mode)
     if console == 'ssh':
-        ssh_client = _connect_to_ssh()
+        ssh_client = ssh_handler()
         # pylint: disable=W0703, C0103
         _, __, return_code = run_ssh_cmd(ssh_client, f"./{script}", timeout=timeout, mode=mode)
         if return_code != 0:
@@ -2491,6 +2511,7 @@ def signal_handler():
     """
 
     print('You pressed Ctrl+C!')
+    close_ssh_connections(SSH_CONNECTIONS)
     kpi.print_kpi_metrics()
     sys.exit(1)
 
@@ -2522,6 +2543,7 @@ def log_heading_msg(msg, pattern='#', panel_size=20):
 
 # pylint: disable=invalid-name
 if __name__ == "__main__":
+    try:
         kpi.init_kpi_metrics()
         signal.signal(signal.SIGINT, signal_handler)
 
@@ -2618,3 +2640,5 @@ if __name__ == "__main__":
 
         LOG.info("INSTALL SUCCEEDED!")
         kpi.print_kpi_metrics()
+    finally:
+        close_ssh_connections(SSH_CONNECTIONS)