resources/libraries/python/ssh.py

   1 # Copyright (c) 2019 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
  17 import socket
  18 import StringIO
  19
  20 from time import time, sleep
  21
  22 from paramiko import RSAKey, SSHClient, AutoAddPolicy
  23 from paramiko.ssh_exception import SSHException, NoValidConnectionsError
  24 from robot.api import logger
  25 from scp import SCPClient, SCPException
  26
  27
  28 __all__ = ["exec_cmd", "exec_cmd_no_error"]
  29
  30 # TODO: load priv key
  31
  32
  33 class SSHTimeout(Exception):
  34     """This exception is raised when a timeout occurs."""
  35     pass
  36
  37
  38 class SSH(object):
  39     """Contains methods for managing and using SSH connections."""
  40
  41     __MAX_RECV_BUF = 10*1024*1024
  42     __existing_connections = {}
  43
  44     def __init__(self):
  45         self._ssh = None
  46         self._node = None
  47
  48     @staticmethod
  49     def _node_hash(node):
  50         """Get IP address and port hash from node dictionary.
  51
  52         :param node: Node in topology.
  53         :type node: dict
  54         :returns: IP address and port for the specified node.
  55         :rtype: int
  56         """
  57
  58         return hash(frozenset([node['host'], node['port']]))
  59
  60     def connect(self, node, attempts=5):
  61         """Connect to node prior to running exec_command or scp.
  62
  63         If there already is a connection to the node, this method reuses it.
  64
  65         :param node: Node in topology.
  66         :param attempts: Number of reconnect attempts.
  67         :type node: dict
  68         :type attempts: int
  69         :raises IOError: If cannot connect to host.
  70         """
  71         self._node = node
  72         node_hash = self._node_hash(node)
  73         if node_hash in SSH.__existing_connections:
  74             self._ssh = SSH.__existing_connections[node_hash]
  75             if self._ssh.get_transport().is_active():
  76                 logger.debug('Reusing SSH: {ssh}'.format(ssh=self._ssh))
  77             else:
  78                 if attempts > 0:
  79                     self._reconnect(attempts-1)
  80                 else:
  81                     raise IOError('Cannot connect to {host}'.
  82                                   format(host=node['host']))
  83         else:
  84             try:
  85                 start = time()
  86                 pkey = None
  87                 if 'priv_key' in node:
  88                     pkey = RSAKey.from_private_key(
  89                         StringIO.StringIO(node['priv_key']))
  90
  91                 self._ssh = SSHClient()
  92                 self._ssh.set_missing_host_key_policy(AutoAddPolicy())
  93
  94                 self._ssh.connect(node['host'], username=node['username'],
  95                                   password=node.get('password'), pkey=pkey,
  96                                   port=node['port'])
  97
  98                 self._ssh.get_transport().set_keepalive(10)
  99
 100                 SSH.__existing_connections[node_hash] = self._ssh
 101                 logger.debug('New SSH to {peer} took {total} seconds: {ssh}'.
 102                              format(
 103                                  peer=self._ssh.get_transport().getpeername(),
 104                                  total=(time() - start),
 105                                  ssh=self._ssh))
 106             except SSHException:
 107                 raise IOError('Cannot connect to {host}'.
 108                               format(host=node['host']))
 109             except NoValidConnectionsError as err:
 110                 logger.error(repr(err))
 111                 raise IOError('Unable to connect to port {port} on {host}'.
 112                               format(port=node['port'], host=node['host']))
 113
 114     def disconnect(self, node=None):
 115         """Close SSH connection to the node.
 116
 117         :param node: The node to disconnect from. None means last connected.
 118         :type node: dict or None
 119         """
 120         if node is None:
 121             node = self._node
 122         if node is None:
 123             return
 124         node_hash = self._node_hash(node)
 125         if node_hash in SSH.__existing_connections:
 126             logger.debug('Disconnecting peer: {host}, {port}'.
 127                          format(host=node['host'], port=node['port']))
 128             ssh = SSH.__existing_connections.pop(node_hash)
 129             ssh.close()
 130
 131     def _reconnect(self, attempts=0):
 132         """Close the SSH connection and open it again.
 133
 134         :param attempts: Number of reconnect attempts.
 135         :type attempts: int
 136         """
 137         node = self._node
 138         self.disconnect(node)
 139         self.connect(node, attempts)
 140         logger.debug('Reconnecting peer done: {host}, {port}'.
 141                      format(host=node['host'], port=node['port']))
 142
 143     def exec_command(self, cmd, timeout=10, log_stdout_err=True):
 144         """Execute SSH command on a new channel on the connected Node.
 145
 146         :param cmd: Command to run on the Node.
 147         :param timeout: Maximal time in seconds to wait until the command is
 148             done. If set to None then wait forever.
 149         :param log_stdout_err: If True, stdout and stderr are logged. stdout
 150             and stderr are logged also if the return code is not zero
 151             independently of the value of log_stdout_err.
 152         :type cmd: str or OptionString
 153         :type timeout: int
 154         :type log_stdout_err: bool
 155         :returns: return_code, stdout, stderr
 156         :rtype: tuple(int, str, str)
 157         :raises SSHTimeout: If command is not finished in timeout time.
 158         """
 159         cmd = str(cmd)
 160         stdout = StringIO.StringIO()
 161         stderr = StringIO.StringIO()
 162         try:
 163             chan = self._ssh.get_transport().open_session(timeout=5)
 164             peer = self._ssh.get_transport().getpeername()
 165         except AttributeError:
 166             self._reconnect()
 167             chan = self._ssh.get_transport().open_session(timeout=5)
 168             peer = self._ssh.get_transport().getpeername()
 169         except SSHException:
 170             self._reconnect()
 171             chan = self._ssh.get_transport().open_session(timeout=5)
 172             peer = self._ssh.get_transport().getpeername()
 173         chan.settimeout(timeout)
 174
 175         logger.trace('exec_command on {peer} with timeout {timeout}: {cmd}'
 176                      .format(peer=peer, timeout=timeout, cmd=cmd))
 177
 178         start = time()
 179         chan.exec_command(cmd)
 180         while not chan.exit_status_ready() and timeout is not None:
 181             if chan.recv_ready():
 182                 stdout.write(chan.recv(self.__MAX_RECV_BUF))
 183
 184             if chan.recv_stderr_ready():
 185                 stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
 186
 187             if time() - start > timeout:
 188                 raise SSHTimeout(
 189                     'Timeout exception during execution of command: {cmd}\n'
 190                     'Current contents of stdout buffer: {stdout}\n'
 191                     'Current contents of stderr buffer: {stderr}\n'
 192                     .format(cmd=cmd, stdout=stdout.getvalue(),
 193                             stderr=stderr.getvalue())
 194                 )
 195
 196             sleep(0.1)
 197         return_code = chan.recv_exit_status()
 198
 199         while chan.recv_ready():
 200             stdout.write(chan.recv(self.__MAX_RECV_BUF))
 201
 202         while chan.recv_stderr_ready():
 203             stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
 204
 205         end = time()
 206         logger.trace('exec_command on {peer} took {total} seconds'.
 207                      format(peer=peer, total=end-start))
 208
 209         logger.trace('return RC {rc}'.format(rc=return_code))
 210         if log_stdout_err or int(return_code):
 211             logger.trace('return STDOUT {stdout}'.
 212                          format(stdout=stdout.getvalue()))
 213             logger.trace('return STDERR {stderr}'.
 214                          format(stderr=stderr.getvalue()))
 215         return return_code, stdout.getvalue(), stderr.getvalue()
 216
 217     def exec_command_sudo(self, cmd, cmd_input=None, timeout=30,
 218                           log_stdout_err=True):
 219         """Execute SSH command with sudo on a new channel on the connected Node.
 220
 221         :param cmd: Command to be executed.
 222         :param cmd_input: Input redirected to the command.
 223         :param timeout: Timeout.
 224         :param log_stdout_err: If True, stdout and stderr are logged.
 225         :type cmd: str
 226         :type cmd_input: str
 227         :type timeout: int
 228         :type log_stdout_err: bool
 229         :returns: return_code, stdout, stderr
 230         :rtype: tuple(int, str, str)
 231
 232         :Example:
 233
 234         >>> from ssh import SSH
 235         >>> ssh = SSH()
 236         >>> ssh.connect(node)
 237         >>> # Execute command without input (sudo -S cmd)
 238         >>> ssh.exec_command_sudo("ifconfig eth0 down")
 239         >>> # Execute command with input (sudo -S cmd <<< "input")
 240         >>> ssh.exec_command_sudo("vpp_api_test", "dump_interface_table")
 241         """
 242         if cmd_input is None:
 243             command = 'sudo -S {c}'.format(c=cmd)
 244         else:
 245             command = 'sudo -S {c} <<< "{i}"'.format(c=cmd, i=cmd_input)
 246         return self.exec_command(command, timeout,
 247                                  log_stdout_err=log_stdout_err)
 248
 249     def exec_command_lxc(self, lxc_cmd, lxc_name, lxc_params='', sudo=True,
 250                          timeout=30):
 251         """Execute command in LXC on a new SSH channel on the connected Node.
 252
 253         :param lxc_cmd: Command to be executed.
 254         :param lxc_name: LXC name.
 255         :param lxc_params: Additional parameters for LXC attach.
 256         :param sudo: Run in privileged LXC mode. Default: privileged
 257         :param timeout: Timeout.
 258         :type lxc_cmd: str
 259         :type lxc_name: str
 260         :type lxc_params: str
 261         :type sudo: bool
 262         :type timeout: int
 263         :returns: return_code, stdout, stderr
 264         """
 265         command = "lxc-attach {p} --name {n} -- /bin/sh -c '{c}'"\
 266             .format(p=lxc_params, n=lxc_name, c=lxc_cmd)
 267
 268         if sudo:
 269             command = 'sudo -S {c}'.format(c=command)
 270         return self.exec_command(command, timeout)
 271
 272     def interactive_terminal_open(self, time_out=45):
 273         """Open interactive terminal on a new channel on the connected Node.
 274
 275         :param time_out: Timeout in seconds.
 276         :returns: SSH channel with opened terminal.
 277
 278         .. warning:: Interruptingcow is used here, and it uses
 279            signal(SIGALRM) to let the operating system interrupt program
 280            execution. This has the following limitations: Python signal
 281            handlers only apply to the main thread, so you cannot use this
 282            from other threads. You must not use this in a program that
 283            uses SIGALRM itself (this includes certain profilers)
 284         """
 285         chan = self._ssh.get_transport().open_session()
 286         chan.get_pty()
 287         chan.invoke_shell()
 288         chan.settimeout(int(time_out))
 289         chan.set_combine_stderr(True)
 290
 291         buf = ''
 292         while not buf.endswith((":~# ", ":~$ ", "~]$ ", "~]# ")):
 293             try:
 294                 chunk = chan.recv(self.__MAX_RECV_BUF)
 295                 if not chunk:
 296                     break
 297                 buf += chunk
 298                 if chan.exit_status_ready():
 299                     logger.error('Channel exit status ready')
 300                     break
 301             except socket.timeout:
 302                 logger.error('Socket timeout: {0}'.format(buf))
 303                 raise Exception('Socket timeout: {0}'.format(buf))
 304         return chan
 305
 306     def interactive_terminal_exec_command(self, chan, cmd, prompt):
 307         """Execute command on interactive terminal.
 308
 309         interactive_terminal_open() method has to be called first!
 310
 311         :param chan: SSH channel with opened terminal.
 312         :param cmd: Command to be executed.
 313         :param prompt: Command prompt, sequence of characters used to
 314         indicate readiness to accept commands.
 315         :returns: Command output.
 316
 317         .. warning:: Interruptingcow is used here, and it uses
 318            signal(SIGALRM) to let the operating system interrupt program
 319            execution. This has the following limitations: Python signal
 320            handlers only apply to the main thread, so you cannot use this
 321            from other threads. You must not use this in a program that
 322            uses SIGALRM itself (this includes certain profilers)
 323         """
 324         chan.sendall('{c}\n'.format(c=cmd))
 325         buf = ''
 326         while not buf.endswith(prompt):
 327             try:
 328                 chunk = chan.recv(self.__MAX_RECV_BUF)
 329                 if not chunk:
 330                     break
 331                 buf += chunk
 332                 if chan.exit_status_ready():
 333                     logger.error('Channel exit status ready')
 334                     break
 335             except socket.timeout:
 336                 logger.error('Socket timeout during execution of command: '
 337                              '{0}\nBuffer content:\n{1}'.format(cmd, buf))
 338                 raise Exception('Socket timeout during execution of command: '
 339                                 '{0}\nBuffer content:\n{1}'.format(cmd, buf))
 340         tmp = buf.replace(cmd.replace('\n', ''), '')
 341         for item in prompt:
 342             tmp.replace(item, '')
 343         return tmp
 344
 345     @staticmethod
 346     def interactive_terminal_close(chan):
 347         """Close interactive terminal SSH channel.
 348
 349         :param chan: SSH channel to be closed.
 350         """
 351         chan.close()
 352
 353     def scp(self, local_path, remote_path, get=False, timeout=30,
 354             wildcard=False):
 355         """Copy files from local_path to remote_path or vice versa.
 356
 357         connect() method has to be called first!
 358
 359         :param local_path: Path to local file that should be uploaded; or
 360         path where to save remote file.
 361         :param remote_path: Remote path where to place uploaded file; or
 362         path to remote file which should be downloaded.
 363         :param get: scp operation to perform. Default is put.
 364         :param timeout: Timeout value in seconds.
 365         :param wildcard: If path has wildcard characters. Default is false.
 366         :type local_path: str
 367         :type remote_path: str
 368         :type get: bool
 369         :type timeout: int
 370         :type wildcard: bool
 371         """
 372         if not get:
 373             logger.trace('SCP {0} to {1}:{2}'.format(
 374                 local_path, self._ssh.get_transport().getpeername(),
 375                 remote_path))
 376         else:
 377             logger.trace('SCP {0}:{1} to {2}'.format(
 378                 self._ssh.get_transport().getpeername(), remote_path,
 379                 local_path))
 380         # SCPCLient takes a paramiko transport as its only argument
 381         if not wildcard:
 382             scp = SCPClient(self._ssh.get_transport(), socket_timeout=timeout)
 383         else:
 384             scp = SCPClient(self._ssh.get_transport(), sanitize=lambda x: x,
 385                             socket_timeout=timeout)
 386         start = time()
 387         if not get:
 388             scp.put(local_path, remote_path)
 389         else:
 390             scp.get(remote_path, local_path)
 391         scp.close()
 392         end = time()
 393         logger.trace('SCP took {0} seconds'.format(end-start))
 394
 395
 396 def exec_cmd(node, cmd, timeout=600, sudo=False, disconnect=False):
 397     """Convenience function to ssh/exec/return rc, out & err.
 398
 399     Returns (rc, stdout, stderr).
 400
 401     :param node: The node to execute command on.
 402     :param cmd: Command to execute.
 403     :param timeout: Timeout value in seconds. Default: 600.
 404     :param sudo: Sudo privilege execution flag. Default: False.
 405     :param disconnect: Close the opened SSH connection if True.
 406     :type node: dict
 407     :type cmd: str or OptionString
 408     :type timeout: int
 409     :type sudo: bool
 410     :type disconnect: bool
 411     :returns: RC, Stdout, Stderr.
 412     :rtype: tuple(int, str, str)
 413     """
 414     if node is None:
 415         raise TypeError('Node parameter is None')
 416     if cmd is None:
 417         raise TypeError('Command parameter is None')
 418     if not cmd:
 419         raise ValueError('Empty command parameter')
 420
 421     ssh = SSH()
 422
 423     if node.get('host_port') is not None:
 424         ssh_node = dict()
 425         ssh_node['host'] = '127.0.0.1'
 426         ssh_node['port'] = node['port']
 427         ssh_node['username'] = node['username']
 428         ssh_node['password'] = node['password']
 429         import pexpect
 430         options = '-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
 431         tnl = '-L {port}:127.0.0.1:{port}'.format(port=node['port'])
 432         ssh_cmd = 'ssh {tnl} {op} {user}@{host} -p {host_port}'.\
 433             format(tnl=tnl, op=options, user=node['host_username'],
 434                    host=node['host'], host_port=node['host_port'])
 435         logger.trace('Initializing local port forwarding:\n{ssh_cmd}'.
 436                      format(ssh_cmd=ssh_cmd))
 437         child = pexpect.spawn(ssh_cmd)
 438         child.expect('.* password: ')
 439         logger.trace(child.after)
 440         child.sendline(node['host_password'])
 441         child.expect('Welcome .*')
 442         logger.trace(child.after)
 443         logger.trace('Local port forwarding finished.')
 444     else:
 445         ssh_node = node
 446
 447     try:
 448         ssh.connect(ssh_node)
 449     except SSHException as err:
 450         logger.error("Failed to connect to node" + repr(err))
 451         return None, None, None
 452
 453     try:
 454         if not sudo:
 455             (ret_code, stdout, stderr) = ssh.exec_command(cmd, timeout=timeout)
 456         else:
 457             (ret_code, stdout, stderr) = ssh.exec_command_sudo(cmd,
 458                                                                timeout=timeout)
 459     except SSHException as err:
 460         logger.error(repr(err))
 461         return None, None, None
 462     finally:
 463         if disconnect:
 464             ssh.disconnect()
 465
 466     return ret_code, stdout, stderr
 467
 468
 469 def exec_cmd_no_error(
 470         node, cmd, timeout=600, sudo=False, message=None, disconnect=False,
 471         retries=0):
 472     """Convenience function to ssh/exec/return out & err.
 473
 474     Verifies that return code is zero.
 475     Supports retries, timeout is related to each try separately then. There is
 476     sleep(1) before each retry.
 477     Disconnect (if enabled) is applied after each try.
 478
 479     :param node: DUT node.
 480     :param cmd: Command to be executed.
 481     :param timeout: Timeout value in seconds. Default: 600.
 482     :param sudo: Sudo privilege execution flag. Default: False.
 483     :param message: Error message in case of failure. Default: None.
 484     :param disconnect: Close the opened SSH connection if True.
 485     :param retries: How many times to retry on failure.
 486     :type node: dict
 487     :type cmd: str or OptionString
 488     :type timeout: int
 489     :type sudo: bool
 490     :type message: str
 491     :type disconnect: bool
 492     :type retries: int
 493     :returns: Stdout, Stderr.
 494     :rtype: tuple(str, str)
 495     :raises RuntimeError: If bash return code is not 0.
 496     """
 497     for _ in range(retries + 1):
 498         ret_code, stdout, stderr = exec_cmd(
 499             node, cmd, timeout=timeout, sudo=sudo, disconnect=disconnect)
 500         if ret_code == 0:
 501             break
 502         sleep(1)
 503     else:
 504         msg = ('Command execution failed: "{cmd}"\n{stderr}'.
 505                format(cmd=cmd, stderr=stderr) if message is None else message)
 506         raise RuntimeError(msg)
 507
 508     return stdout, stderr
 509
 510 def scp_node(
 511         node, local_path, remote_path, get=False, timeout=30, disconnect=False):
 512     """Copy files from local_path to remote_path or vice versa.
 513
 514     :param node: SUT node.
 515     :param local_path: Path to local file that should be uploaded; or
 516         path where to save remote file.
 517     :param remote_path: Remote path where to place uploaded file; or
 518         path to remote file which should be downloaded.
 519     :param get: scp operation to perform. Default is put.
 520     :param timeout: Timeout value in seconds.
 521     :param disconnect: Close the opened SSH connection if True.
 522     :type node: dict
 523     :type local_path: str
 524     :type remote_path: str
 525     :type get: bool
 526     :type timeout: int
 527     :type disconnect: bool
 528     :raises RuntimeError: If SSH connection failed or SCP transfer failed.
 529     """
 530     ssh = SSH()
 531
 532     try:
 533         ssh.connect(node)
 534     except SSHException:
 535         raise RuntimeError('Failed to connect to {host}!'
 536                            .format(host=node['host']))
 537     try:
 538         ssh.scp(local_path, remote_path, get, timeout)
 539     except SCPException:
 540         raise RuntimeError('SCP execution failed on {host}!'
 541                            .format(host=node['host']))
 542     finally:
 543         if disconnect:
 544             ssh.disconnect()