resources/libraries/python/ssh.py

   1 # Copyright (c) 2016 Cisco and/or its affiliates.
   2 # Licensed under the Apache License, Version 2.0 (the "License");
   3 # you may not use this file except in compliance with the License.
   4 # You may obtain a copy of the License at:
   5 #
   6 #     http://www.apache.org/licenses/LICENSE-2.0
   7 #
   8 # Unless required by applicable law or agreed to in writing, software
   9 # distributed under the License is distributed on an "AS IS" BASIS,
  10 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  11 # See the License for the specific language governing permissions and
  12 # limitations under the License.
  13
  14 """Library for SSH connection management."""
  15
  16 import StringIO
  17 from time import time, sleep
  18
  19 import socket
  20 import paramiko
  21 from paramiko import RSAKey
  22 from paramiko.ssh_exception import SSHException
  23 from scp import SCPClient
  24 from robot.api import logger
  25 from robot.utils.asserts import assert_equal
  26
  27 __all__ = ["exec_cmd", "exec_cmd_no_error"]
  28
  29 # TODO: load priv key
  30
  31
  32 class SSH(object):
  33     """Contains methods for managing and using SSH connections."""
  34
  35     __MAX_RECV_BUF = 10*1024*1024
  36     __existing_connections = {}
  37
  38     def __init__(self):
  39         self._ssh = None
  40         self._node = None
  41
  42     @staticmethod
  43     def _node_hash(node):
  44         """Get IP address and port hash from node dictionary.
  45
  46         :param node: Node in topology.
  47         :type node: dict
  48         :return: IP address and port for the specified node.
  49         :rtype: int
  50         """
  51
  52         return hash(frozenset([node['host'], node['port']]))
  53
  54     def connect(self, node):
  55         """Connect to node prior to running exec_command or scp.
  56
  57         If there already is a connection to the node, this method reuses it.
  58         """
  59         self._node = node
  60         node_hash = self._node_hash(node)
  61         if node_hash in SSH.__existing_connections:
  62             self._ssh = SSH.__existing_connections[node_hash]
  63             logger.debug('reusing ssh: {0}'.format(self._ssh))
  64         else:
  65             start = time()
  66             pkey = None
  67             if 'priv_key' in node:
  68                 pkey = RSAKey.from_private_key(
  69                     StringIO.StringIO(node['priv_key']))
  70
  71             self._ssh = paramiko.SSHClient()
  72             self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  73
  74             self._ssh.connect(node['host'], username=node['username'],
  75                               password=node.get('password'), pkey=pkey,
  76                               port=node['port'])
  77
  78             self._ssh.get_transport().set_keepalive(10)
  79
  80             SSH.__existing_connections[node_hash] = self._ssh
  81
  82             logger.trace('connect took {} seconds'.format(time() - start))
  83             logger.debug('new ssh: {0}'.format(self._ssh))
  84
  85         logger.debug('Connect peer: {0}'.
  86                      format(self._ssh.get_transport().getpeername()))
  87         logger.debug('Connections: {0}'.format(str(SSH.__existing_connections)))
  88
  89     def disconnect(self, node):
  90         """Close SSH connection to the node.
  91
  92         :param node: The node to disconnect from.
  93         :type node: dict
  94         """
  95         node_hash = self._node_hash(node)
  96         if node_hash in SSH.__existing_connections:
  97             logger.debug('Disconnecting peer: {}, {}'.
  98                          format(node['host'], node['port']))
  99             ssh = SSH.__existing_connections.pop(node_hash)
 100             ssh.close()
 101
 102     def _reconnect(self):
 103         """Close the SSH connection and open it again."""
 104
 105         node = self._node
 106         self.disconnect(node)
 107         self.connect(node)
 108         logger.debug('Reconnecting peer done: {}'.
 109                      format(self._ssh.get_transport().getpeername()))
 110
 111     def exec_command(self, cmd, timeout=10):
 112         """Execute SSH command on a new channel on the connected Node.
 113
 114         :param cmd: Command to run on the Node.
 115         :param timeout: Maximal time in seconds to wait until the command is
 116         done. If set to None then wait forever.
 117         :type cmd: str
 118         :type timeout: int
 119         :return return_code, stdout, stderr
 120         :rtype: tuple(int, str, str)
 121         :raise socket.timeout: If command is not finished in timeout time.
 122         """
 123         start = time()
 124         stdout = StringIO.StringIO()
 125         stderr = StringIO.StringIO()
 126         try:
 127             chan = self._ssh.get_transport().open_session(timeout=5)
 128         except AttributeError:
 129             self._reconnect()
 130             chan = self._ssh.get_transport().open_session(timeout=5)
 131         except SSHException:
 132             self._reconnect()
 133             chan = self._ssh.get_transport().open_session(timeout=5)
 134         chan.settimeout(timeout)
 135         logger.trace('exec_command on {0}: {1}'
 136                      .format(self._ssh.get_transport().getpeername(), cmd))
 137
 138         chan.exec_command(cmd)
 139         while not chan.exit_status_ready() and timeout is not None:
 140             if chan.recv_ready():
 141                 stdout.write(chan.recv(self.__MAX_RECV_BUF))
 142
 143             if chan.recv_stderr_ready():
 144                 stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
 145
 146             if time() - start > timeout:
 147                 raise socket.timeout(
 148                     'Timeout exception.\n'
 149                     'Current contents of stdout buffer: {0}\n'
 150                     'Current contents of stderr buffer: {1}\n'
 151                     .format(stdout.getvalue(), stderr.getvalue())
 152                 )
 153
 154             sleep(0.1)
 155         return_code = chan.recv_exit_status()
 156
 157         while chan.recv_ready():
 158             stdout.write(chan.recv(self.__MAX_RECV_BUF))
 159
 160         while chan.recv_stderr_ready():
 161             stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
 162
 163         end = time()
 164         logger.trace('exec_command on {0} took {1} seconds'.format(
 165             self._ssh.get_transport().getpeername(), end-start))
 166
 167         logger.trace('chan_recv/_stderr took {} seconds'.format(time()-end))
 168
 169         logger.trace('return RC {}'.format(return_code))
 170         logger.trace('return STDOUT {}'.format(stdout.getvalue()))
 171         logger.trace('return STDERR {}'.format(stderr.getvalue()))
 172         return return_code, stdout.getvalue(), stderr.getvalue()
 173
 174     def exec_command_sudo(self, cmd, cmd_input=None, timeout=30):
 175         """Execute SSH command with sudo on a new channel on the connected Node.
 176
 177         :param cmd: Command to be executed.
 178         :param cmd_input: Input redirected to the command.
 179         :param timeout: Timeout.
 180         :return: return_code, stdout, stderr
 181
 182         :Example:
 183
 184         >>> from ssh import SSH
 185         >>> ssh = SSH()
 186         >>> ssh.connect(node)
 187         >>> # Execute command without input (sudo -S cmd)
 188         >>> ssh.exec_command_sudo("ifconfig eth0 down")
 189         >>> # Execute command with input (sudo -S cmd <<< "input")
 190         >>> ssh.exec_command_sudo("vpp_api_test", "dump_interface_table")
 191         """
 192         if cmd_input is None:
 193             command = 'sudo -S {c}'.format(c=cmd)
 194         else:
 195             command = 'sudo -S {c} <<< "{i}"'.format(c=cmd, i=cmd_input)
 196         return self.exec_command(command, timeout)
 197
 198     def interactive_terminal_open(self, time_out=30):
 199         """Open interactive terminal on a new channel on the connected Node.
 200
 201         :param time_out: Timeout in seconds.
 202         :return: SSH channel with opened terminal.
 203
 204         .. warning:: Interruptingcow is used here, and it uses
 205            signal(SIGALRM) to let the operating system interrupt program
 206            execution. This has the following limitations: Python signal
 207            handlers only apply to the main thread, so you cannot use this
 208            from other threads. You must not use this in a program that
 209            uses SIGALRM itself (this includes certain profilers)
 210         """
 211         chan = self._ssh.get_transport().open_session()
 212         chan.get_pty()
 213         chan.invoke_shell()
 214         chan.settimeout(int(time_out))
 215         chan.set_combine_stderr(True)
 216
 217         buf = ''
 218         while not buf.endswith(':~$ '):
 219             try:
 220                 chunk = chan.recv(self.__MAX_RECV_BUF)
 221                 if not chunk:
 222                     break
 223                 buf += chunk
 224                 if chan.exit_status_ready():
 225                     logger.error('Channel exit status ready')
 226                     break
 227             except socket.timeout:
 228                 raise Exception('Socket timeout: {0}'.format(buf))
 229         return chan
 230
 231     def interactive_terminal_exec_command(self, chan, cmd, prompt):
 232         """Execute command on interactive terminal.
 233
 234         interactive_terminal_open() method has to be called first!
 235
 236         :param chan: SSH channel with opened terminal.
 237         :param cmd: Command to be executed.
 238         :param prompt: Command prompt, sequence of characters used to
 239         indicate readiness to accept commands.
 240         :return: Command output.
 241
 242         .. warning:: Interruptingcow is used here, and it uses
 243            signal(SIGALRM) to let the operating system interrupt program
 244            execution. This has the following limitations: Python signal
 245            handlers only apply to the main thread, so you cannot use this
 246            from other threads. You must not use this in a program that
 247            uses SIGALRM itself (this includes certain profilers)
 248         """
 249         chan.sendall('{c}\n'.format(c=cmd))
 250         buf = ''
 251         while not buf.endswith(prompt):
 252             try:
 253                 chunk = chan.recv(self.__MAX_RECV_BUF)
 254                 if not chunk:
 255                     break
 256                 buf += chunk
 257                 if chan.exit_status_ready():
 258                     logger.error('Channel exit status ready')
 259                     break
 260             except socket.timeout:
 261                 raise Exception('Socket timeout: {0}'.format(buf))
 262         tmp = buf.replace(cmd.replace('\n', ''), '')
 263         return tmp.replace(prompt, '')
 264
 265     @staticmethod
 266     def interactive_terminal_close(chan):
 267         """Close interactive terminal SSH channel.
 268
 269         :param: chan: SSH channel to be closed.
 270         """
 271         chan.close()
 272
 273     def scp(self, local_path, remote_path):
 274         """Copy files from local_path to remote_path.
 275
 276         connect() method has to be called first!
 277         """
 278         logger.trace('SCP {0} to {1}:{2}'.format(
 279             local_path, self._ssh.get_transport().getpeername(), remote_path))
 280         # SCPCLient takes a paramiko transport as its only argument
 281         scp = SCPClient(self._ssh.get_transport())
 282         start = time()
 283         scp.put(local_path, remote_path)
 284         scp.close()
 285         end = time()
 286         logger.trace('SCP took {0} seconds'.format(end-start))
 287
 288
 289 def exec_cmd(node, cmd, timeout=600, sudo=False):
 290     """Convenience function to ssh/exec/return rc, out & err.
 291
 292     Returns (rc, stdout, stderr).
 293     """
 294     if node is None:
 295         raise TypeError('Node parameter is None')
 296     if cmd is None:
 297         raise TypeError('Command parameter is None')
 298     if len(cmd) == 0:
 299         raise ValueError('Empty command parameter')
 300
 301     ssh = SSH()
 302     try:
 303         ssh.connect(node)
 304     except Exception as err:
 305         logger.error("Failed to connect to node" + str(err))
 306         return None, None, None
 307
 308     try:
 309         if not sudo:
 310             (ret_code, stdout, stderr) = ssh.exec_command(cmd, timeout=timeout)
 311         else:
 312             (ret_code, stdout, stderr) = ssh.exec_command_sudo(cmd,
 313                                                                timeout=timeout)
 314     except Exception as err:
 315         logger.error(err)
 316         return None, None, None
 317
 318     return ret_code, stdout, stderr
 319
 320
 321 def exec_cmd_no_error(node, cmd, timeout=600, sudo=False):
 322     """Convenience function to ssh/exec/return out & err.
 323
 324     Verifies that return code is zero.
 325
 326     Returns (stdout, stderr).
 327     """
 328     (ret_code, stdout, stderr) = exec_cmd(node, cmd, timeout=timeout, sudo=sudo)
 329     assert_equal(ret_code, 0, 'Command execution failed: "{}"\n{}'.
 330                  format(cmd, stderr))
 331     return stdout, stderr