1 # Copyright (c) 2022 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
30 u"exec_cmd", u"exec_cmd_no_error", u"SSH", u"SSHTimeout", u"scp_node"
36 class SSHTimeout(Exception):
37 """This exception is raised when a timeout occurs."""
41 """Contains methods for managing and using SSH connections."""
43 __MAX_RECV_BUF = 10 * 1024 * 1024
44 __existing_connections = dict()
52 """Get IP address and port hash from node dictionary.
54 :param node: Node in topology.
56 :returns: IP address and port for the specified node.
59 return hash(frozenset([node[u"host"], node[u"port"]]))
61 def connect(self, node, attempts=5):
62 """Connect to node prior to running exec_command or scp.
64 If there already is a connection to the node, this method reuses it.
66 :param node: Node in topology.
67 :param attempts: Number of reconnect attempts.
70 :raises IOError: If cannot connect to host.
73 node_hash = self._node_hash(node)
74 if node_hash in SSH.__existing_connections:
75 self._ssh = SSH.__existing_connections[node_hash]
76 if self._ssh.get_transport().is_active():
77 logger.debug(f"Reusing SSH: {self._ssh}")
80 self._reconnect(attempts-1)
82 raise IOError(f"Cannot connect to {node['host']}")
87 if u"priv_key" in node:
88 pkey = RSAKey.from_private_key(StringIO(node[u"priv_key"]))
90 self._ssh = SSHClient()
91 self._ssh.set_missing_host_key_policy(AutoAddPolicy())
94 node[u"host"], username=node[u"username"],
95 password=node.get(u"password"), pkey=pkey,
99 self._ssh.get_transport().set_keepalive(10)
101 SSH.__existing_connections[node_hash] = self._ssh
103 f"New SSH to {self._ssh.get_transport().getpeername()} "
104 f"took {monotonic() - start} seconds: {self._ssh}"
106 except SSHException as exc:
107 raise IOError(f"Cannot connect to {node[u'host']}") from exc
108 except NoValidConnectionsError as err:
110 f"Unable to connect to port {node[u'port']} on "
114 def disconnect(self, node=None):
115 """Close SSH connection to the node.
117 :param node: The node to disconnect from. None means last connected.
118 :type node: dict or None
124 node_hash = self._node_hash(node)
125 if node_hash in SSH.__existing_connections:
127 f"Disconnecting peer: {node[u'host']}, {node[u'port']}"
129 ssh = SSH.__existing_connections.pop(node_hash)
132 def _reconnect(self, attempts=0):
133 """Close the SSH connection and open it again.
135 :param attempts: Number of reconnect attempts.
139 self.disconnect(node)
140 self.connect(node, attempts)
142 f"Reconnecting peer done: {node[u'host']}, {node[u'port']}"
145 def exec_command(self, cmd, timeout=10, log_stdout_err=True):
146 """Execute SSH command on a new channel on the connected Node.
148 :param cmd: Command to run on the Node.
149 :param timeout: Maximal time in seconds to wait until the command is
150 done. If set to None then wait forever.
151 :param log_stdout_err: If True, stdout and stderr are logged. stdout
152 and stderr are logged also if the return code is not zero
153 independently of the value of log_stdout_err.
154 Needed for calls outside Robot (e.g. from reservation script).
155 :type cmd: str or OptionString
157 :type log_stdout_err: bool
158 :returns: return_code, stdout, stderr
159 :rtype: tuple(int, str, str)
160 :raises SSHTimeout: If command is not finished in timeout time.
162 if isinstance(cmd, (list, tuple)):
163 cmd = OptionString(cmd)
168 chan = self._ssh.get_transport().open_session(timeout=5)
169 peer = self._ssh.get_transport().getpeername()
170 except (AttributeError, SSHException):
172 chan = self._ssh.get_transport().open_session(timeout=5)
173 peer = self._ssh.get_transport().getpeername()
174 chan.settimeout(timeout)
176 logger.trace(f"exec_command on {peer} with timeout {timeout}: {cmd}")
179 chan.exec_command(cmd)
180 while not chan.exit_status_ready() and timeout is not None:
181 if chan.recv_ready():
182 s_out = chan.recv(self.__MAX_RECV_BUF)
183 stdout += s_out.decode(encoding=u'utf-8', errors=u'ignore') \
184 if isinstance(s_out, bytes) else s_out
186 if chan.recv_stderr_ready():
187 s_err = chan.recv_stderr(self.__MAX_RECV_BUF)
188 stderr += s_err.decode(encoding=u'utf-8', errors=u'ignore') \
189 if isinstance(s_err, bytes) else s_err
191 duration = monotonic() - start
192 if duration > timeout:
194 f"Timeout exception during execution of command: {cmd}\n"
195 f"Current contents of stdout buffer: "
197 f"Current contents of stderr buffer: "
202 return_code = chan.recv_exit_status()
204 while chan.recv_ready():
205 s_out = chan.recv(self.__MAX_RECV_BUF)
206 stdout += s_out.decode(encoding=u'utf-8', errors=u'ignore') \
207 if isinstance(s_out, bytes) else s_out
209 while chan.recv_stderr_ready():
210 s_err = chan.recv_stderr(self.__MAX_RECV_BUF)
211 stderr += s_err.decode(encoding=u'utf-8', errors=u'ignore') \
212 if isinstance(s_err, bytes) else s_err
214 duration = monotonic() - start
215 logger.trace(f"exec_command on {peer} took {duration} seconds")
217 logger.trace(f"return RC {return_code}")
218 if log_stdout_err or int(return_code):
220 f"return STDOUT {stdout}"
223 f"return STDERR {stderr}"
225 return return_code, stdout, stderr
227 def exec_command_sudo(
228 self, cmd, cmd_input=None, timeout=30, log_stdout_err=True):
229 """Execute SSH command with sudo on a new channel on the connected Node.
231 :param cmd: Command to be executed.
232 :param cmd_input: Input redirected to the command.
233 :param timeout: Timeout.
234 :param log_stdout_err: If True, stdout and stderr are logged.
235 Needed for calls outside Robot (e.g. from reservation script).
239 :type log_stdout_err: bool
240 :returns: return_code, stdout, stderr
241 :rtype: tuple(int, str, str)
245 >>> from ssh import SSH
247 >>> ssh.connect(node)
248 >>> # Execute command without input (sudo -S cmd)
249 >>> ssh.exec_command_sudo(u"ifconfig eth0 down")
250 >>> # Execute command with input (sudo -S cmd <<< 'input')
251 >>> ssh.exec_command_sudo(u"vpp_api_test", u"dump_interface_table")
253 if isinstance(cmd, (list, tuple)):
254 cmd = OptionString(cmd)
255 if cmd_input is None:
256 command = f"sudo -E -S {cmd}"
258 command = f"sudo -E -S {cmd} <<< \"{cmd_input}\""
259 return self.exec_command(
260 command, timeout, log_stdout_err=log_stdout_err
263 def exec_command_lxc(
264 self, lxc_cmd, lxc_name, lxc_params=u"", sudo=True, timeout=30):
265 """Execute command in LXC on a new SSH channel on the connected Node.
267 :param lxc_cmd: Command to be executed.
268 :param lxc_name: LXC name.
269 :param lxc_params: Additional parameters for LXC attach.
270 :param sudo: Run in privileged LXC mode. Default: privileged
271 :param timeout: Timeout.
274 :type lxc_params: str
277 :returns: return_code, stdout, stderr
279 command = f"lxc-attach {lxc_params} --name {lxc_name} -- /bin/sh " \
283 command = f"sudo -E -S {command}"
284 return self.exec_command(command, timeout)
286 def interactive_terminal_open(self, time_out=45):
287 """Open interactive terminal on a new channel on the connected Node.
289 :param time_out: Timeout in seconds.
290 :returns: SSH channel with opened terminal.
292 .. warning:: Interruptingcow is used here, and it uses
293 signal(SIGALRM) to let the operating system interrupt program
294 execution. This has the following limitations: Python signal
295 handlers only apply to the main thread, so you cannot use this
296 from other threads. You must not use this in a program that
297 uses SIGALRM itself (this includes certain profilers)
299 chan = self._ssh.get_transport().open_session()
302 chan.settimeout(int(time_out))
303 chan.set_combine_stderr(True)
306 while not buf.endswith((u":~# ", u":~$ ", u"~]$ ", u"~]# ")):
308 s_out = chan.recv(self.__MAX_RECV_BUF)
311 buf += s_out.decode(encoding=u'utf-8', errors=u'ignore') \
312 if isinstance(s_out, bytes) else s_out
313 if chan.exit_status_ready():
314 logger.error(u"Channel exit status ready")
316 except socket.timeout as exc:
317 raise Exception(f"Socket timeout: {buf}") from exc
320 def interactive_terminal_exec_command(self, chan, cmd, prompt):
321 """Execute command on interactive terminal.
323 interactive_terminal_open() method has to be called first!
325 :param chan: SSH channel with opened terminal.
326 :param cmd: Command to be executed.
327 :param prompt: Command prompt, sequence of characters used to
328 indicate readiness to accept commands.
329 :returns: Command output.
331 .. warning:: Interruptingcow is used here, and it uses
332 signal(SIGALRM) to let the operating system interrupt program
333 execution. This has the following limitations: Python signal
334 handlers only apply to the main thread, so you cannot use this
335 from other threads. You must not use this in a program that
336 uses SIGALRM itself (this includes certain profilers)
338 chan.sendall(f"{cmd}\n")
340 while not buf.endswith(prompt):
342 s_out = chan.recv(self.__MAX_RECV_BUF)
345 buf += s_out.decode(encoding=u'utf-8', errors=u'ignore') \
346 if isinstance(s_out, bytes) else s_out
347 if chan.exit_status_ready():
348 logger.error(u"Channel exit status ready")
350 except socket.timeout as exc:
352 f"Socket timeout during execution of command: {cmd}\n"
353 f"Buffer content:\n{buf}"
355 tmp = buf.replace(cmd.replace(u"\n", u""), u"")
357 tmp.replace(item, u"")
361 def interactive_terminal_close(chan):
362 """Close interactive terminal SSH channel.
364 :param chan: SSH channel to be closed.
369 self, local_path, remote_path, get=False, timeout=30,
371 """Copy files from local_path to remote_path or vice versa.
373 connect() method has to be called first!
375 :param local_path: Path to local file that should be uploaded; or
376 path where to save remote file.
377 :param remote_path: Remote path where to place uploaded file; or
378 path to remote file which should be downloaded.
379 :param get: scp operation to perform. Default is put.
380 :param timeout: Timeout value in seconds.
381 :param wildcard: If path has wildcard characters. Default is false.
382 :type local_path: str
383 :type remote_path: str
390 f"SCP {local_path} to "
391 f"{self._ssh.get_transport().getpeername()}:{remote_path}"
395 f"SCP {self._ssh.get_transport().getpeername()}:{remote_path} "
398 # SCPCLient takes a paramiko transport as its only argument
400 scp = SCPClient(self._ssh.get_transport(), socket_timeout=timeout)
403 self._ssh.get_transport(), sanitize=lambda x: x,
404 socket_timeout=timeout
408 scp.put(local_path, remote_path)
410 scp.get(remote_path, local_path)
412 duration = monotonic() - start
413 logger.trace(f"SCP took {duration} seconds")
417 node, cmd, timeout=600, sudo=False, disconnect=False,
420 """Convenience function to ssh/exec/return rc, out & err.
422 Returns (rc, stdout, stderr).
424 :param node: The node to execute command on.
425 :param cmd: Command to execute.
426 :param timeout: Timeout value in seconds. Default: 600.
427 :param sudo: Sudo privilege execution flag. Default: False.
428 :param disconnect: Close the opened SSH connection if True.
429 :param log_stdout_err: If True, stdout and stderr are logged. stdout
430 and stderr are logged also if the return code is not zero
431 independently of the value of log_stdout_err.
432 Needed for calls outside Robot (e.g. from reservation script).
434 :type cmd: str or OptionString
437 :type disconnect: bool
438 :type log_stdout_err: bool
439 :returns: RC, Stdout, Stderr.
440 :rtype: Tuple[int, str, str]
443 raise TypeError(u"Node parameter is None")
445 raise TypeError(u"Command parameter is None")
447 raise ValueError(u"Empty command parameter")
453 except SSHException as err:
454 logger.error(f"Failed to connect to node {node[u'host']}\n{err!r}")
455 return None, None, None
459 ret_code, stdout, stderr = ssh.exec_command(
460 cmd, timeout=timeout, log_stdout_err=log_stdout_err
463 ret_code, stdout, stderr = ssh.exec_command_sudo(
464 cmd, timeout=timeout, log_stdout_err=log_stdout_err
466 except SSHException as err:
467 logger.error(repr(err))
468 return None, None, None
473 return ret_code, stdout, stderr
476 def exec_cmd_no_error(
477 node, cmd, timeout=600, sudo=False, message=None, disconnect=False,
478 retries=0, include_reason=False, log_stdout_err=True
480 """Convenience function to ssh/exec/return out & err.
482 Verifies that return code is zero.
483 Supports retries, timeout is related to each try separately then. There is
484 sleep(1) before each retry.
485 Disconnect (if enabled) is applied after each try.
487 :param node: DUT node.
488 :param cmd: Command to be executed.
489 :param timeout: Timeout value in seconds. Default: 600.
490 :param sudo: Sudo privilege execution flag. Default: False.
491 :param message: Error message in case of failure. Default: None.
492 :param disconnect: Close the opened SSH connection if True.
493 :param retries: How many times to retry on failure.
494 :param include_reason: Whether default info should be appended to message.
495 :param log_stdout_err: If True, stdout and stderr are logged. stdout
496 and stderr are logged also if the return code is not zero
497 independently of the value of log_stdout_err.
498 Needed for calls outside Robot thread (e.g. parallel framework setup).
500 :type cmd: str or OptionString
504 :type disconnect: bool
506 :type include_reason: bool
507 :type log_stdout_err: bool
508 :returns: Stdout, Stderr.
509 :rtype: tuple(str, str)
510 :raises RuntimeError: If bash return code is not 0.
512 for _ in range(retries + 1):
513 ret_code, stdout, stderr = exec_cmd(
514 node, cmd, timeout=timeout, sudo=sudo, disconnect=disconnect,
515 log_stdout_err=log_stdout_err
521 msg = f"Command execution failed: '{cmd}'\nRC: {ret_code}\n{stderr}"
524 msg = f"{message}\n{msg}" if include_reason else message
525 raise RuntimeError(msg)
527 return stdout, stderr
531 node, local_path, remote_path, get=False, timeout=30, disconnect=False):
532 """Copy files from local_path to remote_path or vice versa.
534 :param node: SUT node.
535 :param local_path: Path to local file that should be uploaded; or
536 path where to save remote file.
537 :param remote_path: Remote path where to place uploaded file; or
538 path to remote file which should be downloaded.
539 :param get: scp operation to perform. Default is put.
540 :param timeout: Timeout value in seconds.
541 :param disconnect: Close the opened SSH connection if True.
543 :type local_path: str
544 :type remote_path: str
547 :type disconnect: bool
548 :raises RuntimeError: If SSH connection failed or SCP transfer failed.
554 except SSHException as exc:
555 raise RuntimeError(f"Failed to connect to {node[u'host']}!") from exc
557 ssh.scp(local_path, remote_path, get, timeout)
558 except SCPException as exc:
559 raise RuntimeError(f"SCP execution failed on {node[u'host']}!") from exc