resources/libraries/python/ssh.py

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