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:
6 # http://www.apache.org/licenses/LICENSE-2.0
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.
14 """Library for SSH connection management."""
19 from io import StringIO
20 from time import monotonic, sleep
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
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
33 u"exec_cmd", u"exec_cmd_no_error", u"SSH", u"SSHTimeout", u"scp_node"
39 class SSHTimeout(Exception):
40 """This exception is raised when a timeout occurs."""
44 """Contains methods for managing and using SSH connections."""
46 __MAX_RECV_BUF = 10 * 1024 * 1024
47 __existing_connections = dict()
55 """Get IP address and port hash from node dictionary.
57 :param node: Node in topology.
59 :returns: IP address and port for the specified node.
62 return hash(frozenset([node[u"host"], node[u"port"]]))
64 def connect(self, node, attempts=5):
65 """Connect to node prior to running exec_command or scp.
67 If there already is a connection to the node, this method reuses it.
69 :param node: Node in topology.
70 :param attempts: Number of reconnect attempts.
73 :raises IOError: If cannot connect to host.
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}")
83 self._reconnect(attempts-1)
85 raise IOError(f"Cannot connect to {node['host']}")
90 if u"priv_key" in node:
91 pkey = RSAKey.from_private_key(StringIO(node[u"priv_key"]))
93 self._ssh = SSHClient()
94 self._ssh.set_missing_host_key_policy(AutoAddPolicy())
97 node[u"host"], username=node[u"username"],
98 password=node.get(u"password"), pkey=pkey,
102 self._ssh.get_transport().set_keepalive(10)
104 SSH.__existing_connections[node_hash] = self._ssh
106 f"New SSH to {self._ssh.get_transport().getpeername()} "
107 f"took {monotonic() - start} seconds: {self._ssh}"
109 except SSHException as exc:
110 raise IOError(f"Cannot connect to {node[u'host']}") from exc
111 except NoValidConnectionsError as err:
113 f"Unable to connect to port {node[u'port']} on "
117 def disconnect(self, node=None):
118 """Close SSH connection to the node.
120 :param node: The node to disconnect from. None means last connected.
121 :type node: dict or None
127 node_hash = self._node_hash(node)
128 if node_hash in SSH.__existing_connections:
130 f"Disconnecting peer: {node[u'host']}, {node[u'port']}"
132 ssh = SSH.__existing_connections.pop(node_hash)
135 def _reconnect(self, attempts=0):
136 """Close the SSH connection and open it again.
138 :param attempts: Number of reconnect attempts.
142 self.disconnect(node)
143 self.connect(node, attempts)
145 f"Reconnecting peer done: {node[u'host']}, {node[u'port']}"
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.
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
161 :type log_stdout_err: bool
163 :returns: return_code, stdout, stderr
164 :rtype: tuple(int, str, str)
165 :raises SSHTimeout: If command is not finished in timeout time.
167 if isinstance(cmd, (list, tuple)):
168 cmd = OptionString(cmd)
173 chan = self._ssh.get_transport().open_session(timeout=5)
174 peer = self._ssh.get_transport().getpeername()
175 except (AttributeError, SSHException):
177 chan = self._ssh.get_transport().open_session(timeout=5)
178 peer = self._ssh.get_transport().getpeername()
179 chan.settimeout(timeout)
181 logger.trace(f"exec_command on {peer} with timeout {timeout}: {cmd}")
184 export_ssh_command(self._node[u"host"], self._node[u"port"], cmd)
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
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
198 duration = monotonic() - start
199 if duration > timeout:
202 host=self._node[u"host"],
203 port=self._node[u"port"],
209 f"Timeout exception during execution of command: {cmd}\n"
210 f"Current contents of stdout buffer: "
212 f"Current contents of stderr buffer: "
217 return_code = chan.recv_exit_status()
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
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
229 duration = monotonic() - start
230 logger.trace(f"exec_command on {peer} took {duration} seconds")
232 logger.trace(f"return RC {return_code}")
233 if log_stdout_err or int(return_code):
235 f"return STDOUT {stdout}"
238 f"return STDERR {stderr}"
242 host=self._node[u"host"],
243 port=self._node[u"port"],
249 return return_code, stdout, stderr
251 def exec_command_sudo(
252 self, cmd, cmd_input=None, timeout=30, log_stdout_err=True,
254 """Execute SSH command with sudo on a new channel on the connected Node.
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).
265 :type log_stdout_err: bool
267 :returns: return_code, stdout, stderr
268 :rtype: tuple(int, str, str)
272 >>> from ssh import 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")
280 if isinstance(cmd, (list, tuple)):
281 cmd = OptionString(cmd)
282 if cmd_input is None:
283 command = f"sudo -E -S {cmd}"
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
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.
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.
301 :type lxc_params: str
304 :returns: return_code, stdout, stderr
306 command = f"lxc-attach {lxc_params} --name {lxc_name} -- /bin/sh " \
310 command = f"sudo -E -S {command}"
311 return self.exec_command(command, timeout)
313 def interactive_terminal_open(self, time_out=45):
314 """Open interactive terminal on a new channel on the connected Node.
316 :param time_out: Timeout in seconds.
317 :returns: SSH channel with opened terminal.
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)
326 chan = self._ssh.get_transport().open_session()
329 chan.settimeout(int(time_out))
330 chan.set_combine_stderr(True)
333 while not buf.endswith((u":~# ", u":~$ ", u"~]$ ", u"~]# ")):
335 s_out = chan.recv(self.__MAX_RECV_BUF)
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")
343 except socket.timeout as exc:
344 raise Exception(f"Socket timeout: {buf}") from exc
347 def interactive_terminal_exec_command(self, chan, cmd, prompt):
348 """Execute command on interactive terminal.
350 interactive_terminal_open() method has to be called first!
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.
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)
365 chan.sendall(f"{cmd}\n")
367 while not buf.endswith(prompt):
369 s_out = chan.recv(self.__MAX_RECV_BUF)
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")
377 except socket.timeout as exc:
379 f"Socket timeout during execution of command: {cmd}\n"
380 f"Buffer content:\n{buf}"
382 tmp = buf.replace(cmd.replace(u"\n", u""), u"")
384 tmp.replace(item, u"")
388 def interactive_terminal_close(chan):
389 """Close interactive terminal SSH channel.
391 :param chan: SSH channel to be closed.
396 self, local_path, remote_path, get=False, timeout=30,
398 """Copy files from local_path to remote_path or vice versa.
400 connect() method has to be called first!
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
417 f"SCP {local_path} to "
418 f"{self._ssh.get_transport().getpeername()}:{remote_path}"
422 f"SCP {self._ssh.get_transport().getpeername()}:{remote_path} "
425 # SCPCLient takes a paramiko transport as its only argument
427 scp = SCPClient(self._ssh.get_transport(), socket_timeout=timeout)
430 self._ssh.get_transport(), sanitize=lambda x: x,
431 socket_timeout=timeout
435 scp.put(local_path, remote_path)
437 scp.get(remote_path, local_path)
439 duration = monotonic() - start
440 logger.trace(f"SCP took {duration} seconds")
444 node, cmd, timeout=600, sudo=False, disconnect=False,
445 log_stdout_err=True, export=True
447 """Convenience function to ssh/exec/return rc, out & err.
449 Returns (rc, stdout, stderr).
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).
462 :type cmd: str or OptionString
465 :type disconnect: bool
466 :type log_stdout_err: bool
468 :returns: RC, Stdout, Stderr.
469 :rtype: Tuple[int, str, str]
472 raise TypeError(u"Node parameter is None")
474 raise TypeError(u"Command parameter is None")
476 raise ValueError(u"Empty command parameter")
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
488 ret_code, stdout, stderr = ssh.exec_command(
489 cmd, timeout=timeout, log_stdout_err=log_stdout_err,
493 ret_code, stdout, stderr = ssh.exec_command_sudo(
494 cmd, timeout=timeout, log_stdout_err=log_stdout_err,
497 except SSHException as err:
498 logger.error(repr(err))
499 return None, None, None
504 return ret_code, stdout, stderr
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
511 """Convenience function to ssh/exec/return out & err.
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.
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).
532 :type cmd: str or OptionString
536 :type disconnect: bool
538 :type include_reason: bool
539 :type log_stdout_err: bool
541 :returns: Stdout, Stderr.
542 :rtype: tuple(str, str)
543 :raises RuntimeError: If bash return code is not 0.
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
554 msg = f"Command execution failed: '{cmd}'\nRC: {ret_code}\n{stderr}"
557 msg = f"{message}\n{msg}" if include_reason else message
558 raise RuntimeError(msg)
560 return stdout, stderr
564 node, local_path, remote_path, get=False, timeout=30, disconnect=False):
565 """Copy files from local_path to remote_path or vice versa.
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.
576 :type local_path: str
577 :type remote_path: str
580 :type disconnect: bool
581 :raises RuntimeError: If SSH connection failed or SCP transfer failed.
587 except SSHException as exc:
588 raise RuntimeError(f"Failed to connect to {node[u'host']}!") from exc
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