CSIT-339: Add Keywords for SMT

[csit.git] / resources / libraries / python / ssh.py

diff --git a/resources/libraries/python/ssh.py b/resources/libraries/python/ssh.py
index 4b3838b..7b15998 100644 (file)
--- a/resources/libraries/python/ssh.py
+++ b/resources/libraries/python/ssh.py
@@ -29,6 +29,11 @@ __all__ = ["exec_cmd", "exec_cmd_no_error"]
 # TODO: load priv key
 
 
+class SSHTimeout(Exception):
+    """This exception is raised when a timeout occurs."""
+    pass
+
+
 class SSH(object):
     """Contains methods for managing and using SSH connections."""
 
@@ -118,7 +123,7 @@ class SSH(object):
         :type timeout: int
         :return return_code, stdout, stderr
         :rtype: tuple(int, str, str)
-        :raise socket.timeout: If command is not finished in timeout time.
+        :raise SSHTimeout: If command is not finished in timeout time.
         """
         start = time()
         stdout = StringIO.StringIO()
@@ -144,7 +149,7 @@ class SSH(object):
                 stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
 
             if time() - start > timeout:
-                raise socket.timeout(
+                raise SSHTimeout(
                     'Timeout exception.\n'
                     'Current contents of stdout buffer: {0}\n'
                     'Current contents of stderr buffer: {1}\n'
@@ -215,7 +220,7 @@ class SSH(object):
         chan.set_combine_stderr(True)
 
         buf = ''
-        while not buf.endswith(':~$ '):
+        while not buf.endswith((":~$ ", "~]$ ")):
             try:
                 chunk = chan.recv(self.__MAX_RECV_BUF)
                 if not chunk:
@@ -260,7 +265,9 @@ class SSH(object):
             except socket.timeout:
                 raise Exception('Socket timeout: {0}'.format(buf))
         tmp = buf.replace(cmd.replace('\n', ''), '')
-        return tmp.replace(prompt, '')
+        for p in prompt:
+            tmp.replace(p, '')
+        return tmp
 
     @staticmethod
     def interactive_terminal_close(chan):
@@ -301,7 +308,7 @@ def exec_cmd(node, cmd, timeout=600, sudo=False):
     ssh = SSH()
     try:
         ssh.connect(node)
-    except Exception as err:
+    except SSHException as err:
         logger.error("Failed to connect to node" + str(err))
         return None, None, None
 
@@ -311,7 +318,7 @@ def exec_cmd(node, cmd, timeout=600, sudo=False):
         else:
             (ret_code, stdout, stderr) = ssh.exec_command_sudo(cmd,
                                                                timeout=timeout)
-    except Exception as err:
+    except SSHException as err:
         logger.error(err)
         return None, None, None