X-Git-Url: https://gerrit.fd.io/r/gitweb?p=csit.git;a=blobdiff_plain;f=resources%2Flibraries%2Fpython%2Fssh.py;h=9c7adc44e14093a4043398d54e02244cb3d96585;hp=4beec5aa83c4979fb6e0efe2addaa62c4e7aa3a9;hb=1469b71f48bb4e397ebaa65fbaf0512be691ec49;hpb=5f809dd2414d50fd6f229fc2ba2ff42d66df5dde

diff --git a/resources/libraries/python/ssh.py b/resources/libraries/python/ssh.py
index 4beec5aa83..9c7adc44e1 100644
--- a/resources/libraries/python/ssh.py
+++ b/resources/libraries/python/ssh.py
@@ -13,15 +13,17 @@
 
 """Library for SSH connection management."""
 
-import paramiko
+
 import socket
 import StringIO
 
-from paramiko import RSAKey
+from time import time, sleep
+
+from paramiko import RSAKey, SSHClient, AutoAddPolicy
 from paramiko.ssh_exception import SSHException, NoValidConnectionsError
 from robot.api import logger
-from scp import SCPClient
-from time import time, sleep
+from scp import SCPClient, SCPException
+
 
 __all__ = ["exec_cmd", "exec_cmd_no_error"]
 
@@ -86,8 +88,8 @@ class SSH(object):
                     pkey = RSAKey.from_private_key(
                         StringIO.StringIO(node['priv_key']))
 
-                self._ssh = paramiko.SSHClient()
-                self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+                self._ssh = SSHClient()
+                self._ssh.set_missing_host_key_policy(AutoAddPolicy())
 
                 self._ssh.connect(node['host'], username=node['username'],
                                   password=node.get('password'), pkey=pkey,
@@ -109,12 +111,16 @@ class SSH(object):
                 raise IOError('Unable to connect to port {port} on {host}'.
                               format(port=node['port'], host=node['host']))
 
-    def disconnect(self, node):
+    def disconnect(self, node=None):
         """Close SSH connection to the node.
 
-        :param node: The node to disconnect from.
-        :type node: dict
+        :param node: The node to disconnect from. None means last connected.
+        :type node: dict or None
         """
+        if node is None:
+            node = self._node
+        if node is None:
+            return
         node_hash = self._node_hash(node)
         if node_hash in SSH.__existing_connections:
             logger.debug('Disconnecting peer: {host}, {port}'.
@@ -140,12 +146,13 @@ class SSH(object):
         :param cmd: Command to run on the Node.
         :param timeout: Maximal time in seconds to wait until the command is
         done. If set to None then wait forever.
-        :type cmd: str
+        :type cmd: str or OptionString
         :type timeout: int
         :return return_code, stdout, stderr
         :rtype: tuple(int, str, str)
         :raise SSHTimeout: If command is not finished in timeout time.
         """
+        cmd = str(cmd)
         stdout = StringIO.StringIO()
         stderr = StringIO.StringIO()
         try:
@@ -371,7 +378,7 @@ class SSH(object):
         logger.trace('SCP took {0} seconds'.format(end-start))
 
 
-def exec_cmd(node, cmd, timeout=600, sudo=False):
+def exec_cmd(node, cmd, timeout=600, sudo=False, disconnect=False):
     """Convenience function to ssh/exec/return rc, out & err.
 
     Returns (rc, stdout, stderr).
@@ -380,10 +387,12 @@ def exec_cmd(node, cmd, timeout=600, sudo=False):
     :param cmd: Command to execute.
     :param timeout: Timeout value in seconds. Default: 600.
     :param sudo: Sudo privilege execution flag. Default: False.
+    :param disconnect: Close the opened SSH connection if True.
     :type node: dict
-    :type cmd: str
+    :type cmd: str or OptionString
     :type timeout: int
     :type sudo: bool
+    :type disconnect: bool
     :returns: RC, Stdout, Stderr.
     :rtype: tuple(int, str, str)
     """
@@ -435,33 +444,86 @@ def exec_cmd(node, cmd, timeout=600, sudo=False):
     except SSHException as err:
         logger.error(repr(err))
         return None, None, None
+    finally:
+        if disconnect:
+            ssh.disconnect()
 
     return ret_code, stdout, stderr
 
 
-def exec_cmd_no_error(node, cmd, timeout=600, sudo=False, message=None):
+def exec_cmd_no_error(
+        node, cmd, timeout=600, sudo=False, message=None, disconnect=False,
+        retries=0):
     """Convenience function to ssh/exec/return out & err.
 
     Verifies that return code is zero.
+    Supports retries, timeout is related to each try separately then. There is
+    sleep(1) before each retry.
+    Disconnect (if enabled) is applied after each try.
 
     :param node: DUT node.
     :param cmd: Command to be executed.
     :param timeout: Timeout value in seconds. Default: 600.
     :param sudo: Sudo privilege execution flag. Default: False.
     :param message: Error message in case of failure. Default: None.
+    :param disconnect: Close the opened SSH connection if True.
+    :param retries: How many times to retry on failure.
     :type node: dict
-    :type cmd: str
+    :type cmd: str or OptionString
     :type timeout: int
     :type sudo: bool
     :type message: str
+    :type disconnect: bool
+    :type retries: int
     :returns: Stdout, Stderr.
     :rtype: tuple(str, str)
     :raises RuntimeError: If bash return code is not 0.
     """
-    ret_code, stdout, stderr = exec_cmd(node, cmd, timeout=timeout, sudo=sudo)
-    msg = ('Command execution failed: "{cmd}"\n{stderr}'.
-           format(cmd=cmd, stderr=stderr) if message is None else message)
-    if ret_code != 0:
+    for _ in range(retries + 1):
+        ret_code, stdout, stderr = exec_cmd(
+            node, cmd, timeout=timeout, sudo=sudo, disconnect=disconnect)
+        if ret_code == 0:
+            break
+        sleep(1)
+    else:
+        msg = ('Command execution failed: "{cmd}"\n{stderr}'.
+               format(cmd=cmd, stderr=stderr) if message is None else message)
         raise RuntimeError(msg)
 
     return stdout, stderr
+
+def scp_node(
+        node, local_path, remote_path, get=False, timeout=30, disconnect=False):
+    """Copy files from local_path to remote_path or vice versa.
+
+    :param node: SUT node.
+    :param local_path: Path to local file that should be uploaded; or
+        path where to save remote file.
+    :param remote_path: Remote path where to place uploaded file; or
+        path to remote file which should be downloaded.
+    :param get: scp operation to perform. Default is put.
+    :param timeout: Timeout value in seconds.
+    :param disconnect: Close the opened SSH connection if True.
+    :type node: dict
+    :type local_path: str
+    :type remote_path: str
+    :type get: bool
+    :type timeout: int
+    :type disconnect: bool
+    :raises RuntimeError: If SSH connection failed or SCP transfer failed.
+    """
+    ssh = SSH()
+
+    try:
+        ssh.connect(node)
+    except SSHException:
+        raise RuntimeError('Failed to connect to {host}!'
+                           .format(host=node['host']))
+    try:
+        ssh.scp(local_path, remote_path, get, timeout)
+    except SCPException:
+        raise RuntimeError('SCP execution failed on {host}!'
+                           .format(host=node['host']))
+    finally:
+        if disconnect:
+            ssh.disconnect()