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 import paramiko
  14 from scp import SCPClient
  15 from time import time
  16 from robot.api import logger
  17 from interruptingcow import timeout
  18 from robot.utils.asserts import assert_equal, assert_not_equal
  19
  20 __all__ = ["exec_cmd", "exec_cmd_no_error"]
  21
  22 # TODO: load priv key
  23
  24
  25 class SSH(object):
  26
  27     __MAX_RECV_BUF = 10*1024*1024
  28     __existing_connections = {}
  29
  30     def __init__(self):
  31         self._ssh = paramiko.SSHClient()
  32         self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  33         self._hostname = None
  34
  35     def _node_hash(self, node):
  36         return hash(frozenset([node['host'], node['port']]))
  37
  38     def connect(self, node):
  39         """Connect to node prior to running exec_command or scp.
  40
  41         If there already is a connection to the node, this method reuses it.
  42         """
  43         self._hostname = node['host']
  44         node_hash = self._node_hash(node)
  45         if node_hash in self.__existing_connections:
  46             self._ssh = self.__existing_connections[node_hash]
  47         else:
  48             start = time()
  49             self._ssh.connect(node['host'], username=node['username'],
  50                               password=node['password'])
  51             self.__existing_connections[node_hash] = self._ssh
  52             logger.trace('connect took {} seconds'.format(time() - start))
  53
  54     def exec_command(self, cmd, timeout=10):
  55         """Execute SSH command on a new channel on the connected Node.
  56
  57         Returns (return_code, stdout, stderr).
  58         """
  59         logger.trace('exec_command on {0}: {1}'.format(self._hostname, cmd))
  60         start = time()
  61         chan = self._ssh.get_transport().open_session()
  62         if timeout is not None:
  63             chan.settimeout(int(timeout))
  64         chan.exec_command(cmd)
  65         end = time()
  66         logger.trace('exec_command "{0}" on {1} took {2} seconds'.format(
  67             cmd, self._hostname, end-start))
  68
  69         stdout = ""
  70         while True:
  71             buf = chan.recv(self.__MAX_RECV_BUF)
  72             stdout += buf
  73             if not buf:
  74                 break
  75
  76         stderr = ""
  77         while True:
  78             buf = chan.recv_stderr(self.__MAX_RECV_BUF)
  79             stderr += buf
  80             if not buf:
  81                 break
  82
  83         return_code = chan.recv_exit_status()
  84         logger.trace('chan_recv/_stderr took {} seconds'.format(time()-end))
  85
  86         return (return_code, stdout, stderr)
  87
  88     def exec_command_sudo(self, cmd, cmd_input=None, timeout=10):
  89         """Execute SSH command with sudo on a new channel on the connected Node.
  90
  91            :param cmd: Command to be executed.
  92            :param cmd_input: Input redirected to the command.
  93            :param timeout: Timeout.
  94            :return: return_code, stdout, stderr
  95
  96            :Example:
  97
  98             >>> from ssh import SSH
  99             >>> ssh = SSH()
 100             >>> ssh.connect(node)
 101             >>> #Execute command without input (sudo -S cmd)
 102             >>> ssh.exex_command_sudo("ifconfig eth0 down")
 103             >>> #Execute command with input (sudo -S cmd <<< "input")
 104             >>> ssh.exex_command_sudo("vpe_api_test", "dump_interface_table")
 105         """
 106         if cmd_input is None:
 107             command = 'sudo -S {c}'.format(c=cmd)
 108         else:
 109             command = 'sudo -S {c} <<< "{i}"'.format(c=cmd, i=cmd_input)
 110         return self.exec_command(command, timeout)
 111
 112     def interactive_terminal_open(self, time_out=10):
 113         """Open interactive terminal on a new channel on the connected Node.
 114
 115            :param time_out: Timeout in seconds.
 116            :return: SSH channel with opened terminal.
 117
 118            .. warning:: Interruptingcow is used here, and it uses
 119                signal(SIGALRM) to let the operating system interrupt program
 120                execution. This has the following limitations: Python signal
 121                handlers only apply to the main thread, so you cannot use this
 122                from other threads. You must not use this in a program that
 123                uses SIGALRM itself (this includes certain profilers)
 124         """
 125         chan = self._ssh.get_transport().open_session()
 126         chan.get_pty()
 127         chan.invoke_shell()
 128         chan.settimeout(int(time_out))
 129
 130         buf = ''
 131         try:
 132             with timeout(time_out, exception=RuntimeError):
 133                 while not buf.endswith(':~$ '):
 134                     if chan.recv_ready():
 135                         buf = chan.recv(4096)
 136         except RuntimeError:
 137             raise Exception('Open interactive terminal timeout.')
 138         return chan
 139
 140     def interactive_terminal_exec_command(self, chan, cmd, prompt,
 141                                           time_out=10):
 142         """Execute command on interactive terminal.
 143
 144            interactive_terminal_open() method has to be called first!
 145
 146            :param chan: SSH channel with opened terminal.
 147            :param cmd: Command to be executed.
 148            :param prompt: Command prompt, sequence of characters used to
 149                indicate readiness to accept commands.
 150            :param time_out: Timeout in seconds.
 151            :return: Command output.
 152
 153            .. warning:: Interruptingcow is used here, and it uses
 154                signal(SIGALRM) to let the operating system interrupt program
 155                execution. This has the following limitations: Python signal
 156                handlers only apply to the main thread, so you cannot use this
 157                from other threads. You must not use this in a program that
 158                uses SIGALRM itself (this includes certain profilers)
 159         """
 160         chan.sendall('{c}\n'.format(c=cmd))
 161         buf = ''
 162         try:
 163             with timeout(time_out, exception=RuntimeError):
 164                 while not buf.endswith(prompt):
 165                     if chan.recv_ready():
 166                         buf += chan.recv(4096)
 167         except RuntimeError:
 168             raise Exception("Exec '{c}' timeout.".format(c=cmd))
 169         tmp = buf.replace(cmd.replace('\n', ''), '')
 170         return tmp.replace(prompt, '')
 171
 172     def interactive_terminal_close(self, chan):
 173         """Close interactive terminal SSH channel.
 174
 175            :param: chan: SSH channel to be closed.
 176         """
 177         chan.close()
 178
 179     def scp(self, local_path, remote_path):
 180         """Copy files from local_path to remote_path.
 181
 182         connect() method has to be called first!
 183         """
 184         logger.trace('SCP {0} to {1}:{2}'.format(
 185             local_path, self._hostname, remote_path))
 186         # SCPCLient takes a paramiko transport as its only argument
 187         scp = SCPClient(self._ssh.get_transport())
 188         start = time()
 189         scp.put(local_path, remote_path)
 190         scp.close()
 191         end = time()
 192         logger.trace('SCP took {0} seconds'.format(end-start))
 193
 194
 195 def exec_cmd(node, cmd, timeout=None, sudo=False):
 196     """Convenience function to ssh/exec/return rc, out & err.
 197
 198     Returns (rc, stdout, stderr).
 199     """
 200     if node is None:
 201         raise TypeError('Node parameter is None')
 202     if cmd is None:
 203         raise TypeError('Command parameter is None')
 204     if len(cmd) == 0:
 205         raise ValueError('Empty command parameter')
 206
 207     ssh = SSH()
 208     try:
 209         ssh.connect(node)
 210     except Exception, e:
 211         logger.error("Failed to connect to node" + e)
 212         return None
 213
 214     try:
 215         if not sudo:
 216             (ret_code, stdout, stderr) = ssh.exec_command(cmd, timeout=timeout)
 217         else:
 218             (ret_code, stdout, stderr) = ssh.exec_command_sudo(cmd,
 219                                                                timeout=timeout)
 220     except Exception, e:
 221         logger.error(e)
 222         return None
 223
 224     return (ret_code, stdout, stderr)
 225
 226 def exec_cmd_no_error(node, cmd, timeout=None, sudo=False):
 227     """Convenience function to ssh/exec/return out & err.
 228     Verifies that return code is zero.
 229
 230     Returns (stdout, stderr).
 231     """
 232     (rc, stdout, stderr) = exec_cmd(node,cmd, timeout=timeout, sudo=sudo)
 233     assert_equal(rc, 0, 'Command execution failed: "{}"\n{}'.
 234                  format(cmd, stderr))
 235     return (stdout, stderr)