Merge "Improve SSH to Open/Close interactive terminal"

2 files changed, 122 insertions, 0 deletions

diff --git a/yardstick/ssh.py b/yardstick/ssh.py
index 2ebf40e98..438e8158b 100644
--- a/yardstick/ssh.py
+++ b/yardstick/ssh.py
@@ -448,6 +448,86 @@ class SSH(object):
         with client.open_sftp() as sftp:
             sftp.getfo(remotepath, file_obj)
 
+    def interactive_terminal_open(self, time_out=45):
+        """Open interactive terminal on a SSH channel.
+
+        :param time_out: Timeout in seconds.
+        :returns: SSH channel with opened terminal.
+
+        .. warning:: Interruptingcow is used here, and it uses
+           signal(SIGALRM) to let the operating system interrupt program
+           execution. This has the following limitations: Python signal
+           handlers only apply to the main thread, so you cannot use this
+           from other threads. You must not use this in a program that
+           uses SIGALRM itself (this includes certain profilers)
+        """
+        chan = self._get_client().get_transport().open_session()
+        chan.get_pty()
+        chan.invoke_shell()
+        chan.settimeout(int(time_out))
+        chan.set_combine_stderr(True)
+
+        buf = ''
+        while not buf.endswith((":~# ", ":~$ ", "~]$ ", "~]# ")):
+            try:
+                chunk = chan.recv(10 * 1024 * 1024)
+                if not chunk:
+                    break
+                buf += chunk
+                if chan.exit_status_ready():
+                    self.log.error('Channel exit status ready')
+                    break
+            except socket.timeout:
+                raise exceptions.SSHTimeout(error_msg='Socket timeout: %s' % buf)
+        return chan
+
+    def interactive_terminal_exec_command(self, chan, cmd, prompt):
+        """Execute command on interactive terminal.
+
+        interactive_terminal_open() method has to be called first!
+
+        :param chan: SSH channel with opened terminal.
+        :param cmd: Command to be executed.
+        :param prompt: Command prompt, sequence of characters used to
+        indicate readiness to accept commands.
+        :returns: Command output.
+
+        .. warning:: Interruptingcow is used here, and it uses
+           signal(SIGALRM) to let the operating system interrupt program
+           execution. This has the following limitations: Python signal
+           handlers only apply to the main thread, so you cannot use this
+           from other threads. You must not use this in a program that
+           uses SIGALRM itself (this includes certain profilers)
+        """
+        chan.sendall('{c}\n'.format(c=cmd))
+        buf = ''
+        while not buf.endswith(prompt):
+            try:
+                chunk = chan.recv(10 * 1024 * 1024)
+                if not chunk:
+                    break
+                buf += chunk
+                if chan.exit_status_ready():
+                    self.log.error('Channel exit status ready')
+                    break
+            except socket.timeout:
+                message = ("Socket timeout during execution of command: "
+                           "%(cmd)s\nBuffer content:\n%(buf)s" % {"cmd": cmd,
+                                                                  "buf": buf})
+                raise exceptions.SSHTimeout(error_msg=message)
+        tmp = buf.replace(cmd.replace('\n', ''), '')
+        for item in prompt:
+            tmp.replace(item, '')
+        return tmp
+
+    @staticmethod
+    def interactive_terminal_close(chan):
+        """Close interactive terminal SSH channel.
+
+        :param: chan: SSH channel to be closed.
+        """
+        chan.close()
+
 
 class AutoConnectSSH(SSH):
 
diff --git a/yardstick/tests/unit/test_ssh.py b/yardstick/tests/unit/test_ssh.py
index 71929f1a2..374fb6644 100644
--- a/yardstick/tests/unit/test_ssh.py
+++ b/yardstick/tests/unit/test_ssh.py
@@ -286,6 +286,48 @@ class SSHTestCase(unittest.TestCase):
         mock_paramiko_exec_command.assert_called_once_with('cmd',
                                                            get_pty=True)
 
+    @mock.patch("yardstick.ssh.paramiko")
+    def test_interactive_terminal_open(self, mock_paramiko):
+        fake_client = mock.Mock()
+        fake_session = mock.Mock()
+        fake_session.recv.return_value = ":~# "
+        fake_transport = mock.Mock()
+        fake_transport.open_session.return_value = fake_session
+        fake_client.get_transport.return_value = fake_transport
+        mock_paramiko.SSHClient.return_value = fake_client
+
+        test_ssh = ssh.SSH("admin", "example.net", pkey="key")
+        result = test_ssh.interactive_terminal_open()
+        self.assertEqual(fake_session, result)
+
+    @mock.patch("yardstick.ssh.paramiko")
+    def test_interactive_terminal_exec_command(self, mock_paramiko):
+        fake_client = mock.Mock()
+        fake_session = mock.Mock()
+        fake_session.recv.return_value = "stdout fake data"
+        fake_transport = mock.Mock()
+        fake_transport.open_session.return_value = fake_session
+        fake_client.get_transport.return_value = fake_transport
+        mock_paramiko.SSHClient.return_value = fake_client
+
+        test_ssh = ssh.SSH("admin", "example.net", pkey="key")
+        with mock.patch.object(fake_session, "sendall") \
+                as mock_paramiko_send_command:
+            result = test_ssh.interactive_terminal_exec_command(fake_session,
+                                                                'cmd', "vat# ")
+        self.assertEqual("stdout fake data", result)
+        mock_paramiko_send_command.assert_called_once_with('cmd\n')
+
+    @mock.patch("yardstick.ssh.paramiko")
+    def test_interactive_terminal_close(self, _):
+        fake_session = mock.Mock()
+        paramiko_sshclient = self.test_client._get_client()
+        paramiko_sshclient.get_transport.open_session.return_value = fake_session
+        with mock.patch.object(fake_session, "close") \
+                as mock_paramiko_terminal_close:
+            self.test_client.interactive_terminal_close(fake_session)
+        mock_paramiko_terminal_close.assert_called_once_with()
+
 
 class SSHRunTestCase(unittest.TestCase):
     """Test SSH.run method in different aspects.