469d77ffc2cef5cba6936d670fa4b9c687007eb2
[csit.git] / 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 SSHTimeout(Exception):
33     """This exception is raised when a timeout occurs."""
34     pass
35
36
37 class SSH(object):
38     """Contains methods for managing and using SSH connections."""
39
40     __MAX_RECV_BUF = 10*1024*1024
41     __existing_connections = {}
42
43     def __init__(self):
44         self._ssh = None
45         self._node = None
46
47     @staticmethod
48     def _node_hash(node):
49         """Get IP address and port hash from node dictionary.
50
51         :param node: Node in topology.
52         :type node: dict
53         :return: IP address and port for the specified node.
54         :rtype: int
55         """
56
57         return hash(frozenset([node['host'], node['port']]))
58
59     def connect(self, node, attempts=5):
60         """Connect to node prior to running exec_command or scp.
61
62         If there already is a connection to the node, this method reuses it.
63         """
64         try:
65             self._node = node
66             node_hash = self._node_hash(node)
67             if node_hash in SSH.__existing_connections:
68                 self._ssh = SSH.__existing_connections[node_hash]
69                 logger.debug('reusing ssh: {0}'.format(self._ssh))
70             else:
71                 start = time()
72                 pkey = None
73                 if 'priv_key' in node:
74                     pkey = RSAKey.from_private_key(
75                         StringIO.StringIO(node['priv_key']))
76
77                 self._ssh = paramiko.SSHClient()
78                 self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
79
80                 self._ssh.connect(node['host'], username=node['username'],
81                                   password=node.get('password'), pkey=pkey,
82                                   port=node['port'])
83
84                 self._ssh.get_transport().set_keepalive(10)
85
86                 SSH.__existing_connections[node_hash] = self._ssh
87
88                 logger.trace('connect took {} seconds'.format(time() - start))
89                 logger.debug('new ssh: {0}'.format(self._ssh))
90
91             logger.debug('Connect peer: {0}'.
92                          format(self._ssh.get_transport().getpeername()))
93             logger.debug('Connections: {0}'.
94                          format(str(SSH.__existing_connections)))
95         except:
96             if attempts > 0:
97                 self._reconnect(attempts-1)
98             else:
99                 raise
100
101     def disconnect(self, node):
102         """Close SSH connection to the node.
103
104         :param node: The node to disconnect from.
105         :type node: dict
106         """
107         node_hash = self._node_hash(node)
108         if node_hash in SSH.__existing_connections:
109             logger.debug('Disconnecting peer: {}, {}'.
110                          format(node['host'], node['port']))
111             ssh = SSH.__existing_connections.pop(node_hash)
112             ssh.close()
113
114     def _reconnect(self, attempts=0):
115         """Close the SSH connection and open it again."""
116
117         node = self._node
118         self.disconnect(node)
119         self.connect(node, attempts)
120         logger.debug('Reconnecting peer done: {}'.
121                      format(self._ssh.get_transport().getpeername()))
122
123     def exec_command(self, cmd, timeout=10):
124         """Execute SSH command on a new channel on the connected Node.
125
126         :param cmd: Command to run on the Node.
127         :param timeout: Maximal time in seconds to wait until the command is
128         done. If set to None then wait forever.
129         :type cmd: str
130         :type timeout: int
131         :return return_code, stdout, stderr
132         :rtype: tuple(int, str, str)
133         :raise SSHTimeout: If command is not finished in timeout time.
134         """
135         stdout = StringIO.StringIO()
136         stderr = StringIO.StringIO()
137         try:
138             chan = self._ssh.get_transport().open_session(timeout=5)
139             peer = self._ssh.get_transport().getpeername()
140         except AttributeError:
141             self._reconnect()
142             chan = self._ssh.get_transport().open_session(timeout=5)
143             peer = self._ssh.get_transport().getpeername()
144         except SSHException:
145             self._reconnect()
146             chan = self._ssh.get_transport().open_session(timeout=5)
147             peer = self._ssh.get_transport().getpeername()
148         chan.settimeout(timeout)
149
150         logger.trace('exec_command on {peer} with timeout {timeout}: {cmd}'
151                      .format(peer=peer, timeout=timeout, cmd=cmd))
152
153         start = time()
154         chan.exec_command(cmd)
155         while not chan.exit_status_ready() and timeout is not None:
156             if chan.recv_ready():
157                 stdout.write(chan.recv(self.__MAX_RECV_BUF))
158
159             if chan.recv_stderr_ready():
160                 stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
161
162             if time() - start > timeout:
163                 raise SSHTimeout(
164                     'Timeout exception during execution of command: {cmd}\n'
165                     'Current contents of stdout buffer: {stdout}\n'
166                     'Current contents of stderr buffer: {stderr}\n'
167                     .format(cmd=cmd, stdout=stdout.getvalue(),
168                             stderr=stderr.getvalue())
169                 )
170
171             sleep(0.1)
172         return_code = chan.recv_exit_status()
173
174         while chan.recv_ready():
175             stdout.write(chan.recv(self.__MAX_RECV_BUF))
176
177         while chan.recv_stderr_ready():
178             stderr.write(chan.recv_stderr(self.__MAX_RECV_BUF))
179
180         end = time()
181         logger.trace('exec_command on {peer} took {total} seconds'.
182                      format(peer=peer, total=end-start))
183
184         logger.trace('return RC {rc}'.format(rc=return_code))
185         logger.trace('return STDOUT {stdout}'.format(stdout=stdout.getvalue()))
186         logger.trace('return STDERR {stderr}'.format(stderr=stderr.getvalue()))
187         return return_code, stdout.getvalue(), stderr.getvalue()
188
189     def exec_command_sudo(self, cmd, cmd_input=None, timeout=30):
190         """Execute SSH command with sudo on a new channel on the connected Node.
191
192         :param cmd: Command to be executed.
193         :param cmd_input: Input redirected to the command.
194         :param timeout: Timeout.
195         :return: return_code, stdout, stderr
196
197         :Example:
198
199         >>> from ssh import SSH
200         >>> ssh = SSH()
201         >>> ssh.connect(node)
202         >>> # Execute command without input (sudo -S cmd)
203         >>> ssh.exec_command_sudo("ifconfig eth0 down")
204         >>> # Execute command with input (sudo -S cmd <<< "input")
205         >>> ssh.exec_command_sudo("vpp_api_test", "dump_interface_table")
206         """
207         if cmd_input is None:
208             command = 'sudo -S {c}'.format(c=cmd)
209         else:
210             command = 'sudo -S {c} <<< "{i}"'.format(c=cmd, i=cmd_input)
211         return self.exec_command(command, timeout)
212
213     def exec_command_lxc(self, lxc_cmd, lxc_name, lxc_params='', sudo=True,
214                          timeout=30):
215         """Execute command in LXC on a new SSH channel on the connected Node.
216
217         :param lxc_cmd: Command to be executed.
218         :param lxc_name: LXC name.
219         :param lxc_params: Additional parameters for LXC attach.
220         :param sudo: Run in privileged LXC mode. Default: privileged
221         :param timeout: Timeout.
222         :type lxc_cmd: str
223         :type lxc_name: str
224         :type lxc_params: str
225         :type sudo: bool
226         :type timeout: int
227         :return: return_code, stdout, stderr
228         """
229         command = "lxc-attach {p} --name {n} -- /bin/sh -c '{c}'"\
230             .format(p=lxc_params, n=lxc_name, c=lxc_cmd)
231
232         if sudo:
233             command = 'sudo -S {c}'.format(c=command)
234         return self.exec_command(command, timeout)
235
236     def interactive_terminal_open(self, time_out=30):
237         """Open interactive terminal on a new channel on the connected Node.
238
239         :param time_out: Timeout in seconds.
240         :return: SSH channel with opened terminal.
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 = self._ssh.get_transport().open_session()
250         chan.get_pty()
251         chan.invoke_shell()
252         chan.settimeout(int(time_out))
253         chan.set_combine_stderr(True)
254
255         buf = ''
256         while not buf.endswith((":~$ ", "~]$ ", "~]# ")):
257             try:
258                 chunk = chan.recv(self.__MAX_RECV_BUF)
259                 if not chunk:
260                     break
261                 buf += chunk
262                 if chan.exit_status_ready():
263                     logger.error('Channel exit status ready')
264                     break
265             except socket.timeout:
266                 logger.error('Socket timeout: {0}'.format(buf))
267                 raise Exception('Socket timeout: {0}'.format(buf))
268         return chan
269
270     def interactive_terminal_exec_command(self, chan, cmd, prompt):
271         """Execute command on interactive terminal.
272
273         interactive_terminal_open() method has to be called first!
274
275         :param chan: SSH channel with opened terminal.
276         :param cmd: Command to be executed.
277         :param prompt: Command prompt, sequence of characters used to
278         indicate readiness to accept commands.
279         :return: Command output.
280
281         .. warning:: Interruptingcow is used here, and it uses
282            signal(SIGALRM) to let the operating system interrupt program
283            execution. This has the following limitations: Python signal
284            handlers only apply to the main thread, so you cannot use this
285            from other threads. You must not use this in a program that
286            uses SIGALRM itself (this includes certain profilers)
287         """
288         chan.sendall('{c}\n'.format(c=cmd))
289         buf = ''
290         while not buf.endswith(prompt):
291             try:
292                 chunk = chan.recv(self.__MAX_RECV_BUF)
293                 if not chunk:
294                     break
295                 buf += chunk
296                 if chan.exit_status_ready():
297                     logger.error('Channel exit status ready')
298                     break
299             except socket.timeout:
300                 logger.error('Socket timeout during execution of command: '
301                              '{0}\nBuffer content:\n{1}'.format(cmd, buf))
302                 raise Exception('Socket timeout during execution of command: '
303                                 '{0}\nBuffer content:\n{1}'.format(cmd, buf))
304         tmp = buf.replace(cmd.replace('\n', ''), '')
305         for item in prompt:
306             tmp.replace(item, '')
307         return tmp
308
309     @staticmethod
310     def interactive_terminal_close(chan):
311         """Close interactive terminal SSH channel.
312
313         :param: chan: SSH channel to be closed.
314         """
315         chan.close()
316
317     def scp(self, local_path, remote_path, get=False, timeout=10):
318         """Copy files from local_path to remote_path or vice versa.
319
320         connect() method has to be called first!
321
322         :param local_path: Path to local file that should be uploaded; or
323         path where to save remote file.
324         :param remote_path: Remote path where to place uploaded file; or
325         path to remote file which should be downloaded.
326         :param get: scp operation to perform. Default is put.
327         :param timeout: Timeout value in seconds.
328         :type local_path: str
329         :type remote_path: str
330         :type get: bool
331         :type timeout: int
332         """
333         if not get:
334             logger.trace('SCP {0} to {1}:{2}'.format(
335                 local_path, self._ssh.get_transport().getpeername(),
336                 remote_path))
337         else:
338             logger.trace('SCP {0}:{1} to {2}'.format(
339                 self._ssh.get_transport().getpeername(), remote_path,
340                 local_path))
341         # SCPCLient takes a paramiko transport as its only argument
342         scp = SCPClient(self._ssh.get_transport(), socket_timeout=timeout)
343         start = time()
344         if not get:
345             scp.put(local_path, remote_path)
346         else:
347             scp.get(remote_path, local_path)
348         scp.close()
349         end = time()
350         logger.trace('SCP took {0} seconds'.format(end-start))
351
352
353 def exec_cmd(node, cmd, timeout=600, sudo=False):
354     """Convenience function to ssh/exec/return rc, out & err.
355
356     Returns (rc, stdout, stderr).
357     """
358     if node is None:
359         raise TypeError('Node parameter is None')
360     if cmd is None:
361         raise TypeError('Command parameter is None')
362     if len(cmd) == 0:
363         raise ValueError('Empty command parameter')
364
365     ssh = SSH()
366     try:
367         ssh.connect(node)
368     except SSHException as err:
369         logger.error("Failed to connect to node" + str(err))
370         return None, None, None
371
372     try:
373         if not sudo:
374             (ret_code, stdout, stderr) = ssh.exec_command(cmd, timeout=timeout)
375         else:
376             (ret_code, stdout, stderr) = ssh.exec_command_sudo(cmd,
377                                                                timeout=timeout)
378     except SSHException as err:
379         logger.error(err)
380         return None, None, None
381
382     return ret_code, stdout, stderr
383
384
385 def exec_cmd_no_error(node, cmd, timeout=600, sudo=False):
386     """Convenience function to ssh/exec/return out & err.
387
388     Verifies that return code is zero.
389
390     Returns (stdout, stderr).
391     """
392     (ret_code, stdout, stderr) = exec_cmd(node, cmd, timeout=timeout, sudo=sudo)
393     assert_equal(ret_code, 0, 'Command execution failed: "{}"\n{}'.
394                  format(cmd, stderr))
395     return stdout, stderr

©2016 FD.io a Linux Foundation Collaborative Project. All Rights Reserved.
Linux Foundation is a registered trademark of The Linux Foundation. Linux is a registered trademark of Linus Torvalds.
Please see our privacy policy and terms of use.