yardstick/ssh.py

   1 # Copyright 2013: Mirantis Inc.
   2 # All Rights Reserved.
   3 #
   4 #    Licensed under the Apache License, Version 2.0 (the "License"); you may
   5 #    not use this file except in compliance with the License. You may obtain
   6 #    a copy of the License at
   7 #
   8 #         http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 #    Unless required by applicable law or agreed to in writing, software
  11 #    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  12 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  13 #    License for the specific language governing permissions and limitations
  14 #    under the License.
  15
  16 # yardstick comment: this is a modified copy of rally/rally/common/sshutils.py
  17
  18 """High level ssh library.
  19
  20 Usage examples:
  21
  22 Execute command and get output:
  23
  24     ssh = sshclient.SSH("root", "example.com", port=33)
  25     status, stdout, stderr = ssh.execute("ps ax")
  26     if status:
  27         raise Exception("Command failed with non-zero status.")
  28     print stdout.splitlines()
  29
  30 Execute command with huge output:
  31
  32     class PseudoFile(object):
  33         def write(chunk):
  34             if "error" in chunk:
  35                 email_admin(chunk)
  36
  37     ssh = sshclient.SSH("root", "example.com")
  38     ssh.run("tail -f /var/log/syslog", stdout=PseudoFile(), timeout=False)
  39
  40 Execute local script on remote side:
  41
  42     ssh = sshclient.SSH("user", "example.com")
  43     status, out, err = ssh.execute("/bin/sh -s arg1 arg2",
  44                                    stdin=open("~/myscript.sh", "r"))
  45
  46 Upload file:
  47
  48     ssh = sshclient.SSH("user", "example.com")
  49     ssh.run("cat > ~/upload/file.gz", stdin=open("/store/file.gz", "rb"))
  50
  51 Eventlet:
  52
  53     eventlet.monkey_patch(select=True, time=True)
  54     or
  55     eventlet.monkey_patch()
  56     or
  57     sshclient = eventlet.import_patched("opentstack.common.sshclient")
  58
  59 """
  60 import os
  61 import select
  62 import socket
  63 import time
  64
  65 import logging
  66 import paramiko
  67 from scp import SCPClient
  68 import six
  69
  70
  71 DEFAULT_PORT = 22
  72
  73
  74 class SSHError(Exception):
  75     pass
  76
  77
  78 class SSHTimeout(SSHError):
  79     pass
  80
  81
  82 class SSH(object):
  83     """Represent ssh connection."""
  84
  85     def __init__(self, user, host, port=DEFAULT_PORT, pkey=None,
  86                  key_filename=None, password=None, name=None):
  87         """Initialize SSH client.
  88
  89         :param user: ssh username
  90         :param host: hostname or ip address of remote ssh server
  91         :param port: remote ssh port
  92         :param pkey: RSA or DSS private key string or file object
  93         :param key_filename: private key filename
  94         :param password: password
  95         """
  96         self.name = name
  97         if name:
  98             self.log = logging.getLogger(__name__ + '.' + self.name)
  99         else:
 100             self.log = logging.getLogger(__name__)
 101
 102         self.user = user
 103         self.host = host
 104         # we may get text port from YAML, convert to int
 105         self.port = int(port)
 106         self.pkey = self._get_pkey(pkey) if pkey else None
 107         self.password = password
 108         self.key_filename = key_filename
 109         self._client = False
 110         # paramiko loglevel debug will output ssh protocl debug
 111         # we don't ever really want that unless we are debugging paramiko
 112         # ssh issues
 113         if os.environ.get("PARAMIKO_DEBUG", "").lower() == "true":
 114             logging.getLogger("paramiko").setLevel(logging.DEBUG)
 115         else:
 116             logging.getLogger("paramiko").setLevel(logging.WARN)
 117
 118     def _get_pkey(self, key):
 119         if isinstance(key, six.string_types):
 120             key = six.moves.StringIO(key)
 121         errors = []
 122         for key_class in (paramiko.rsakey.RSAKey, paramiko.dsskey.DSSKey):
 123             try:
 124                 return key_class.from_private_key(key)
 125             except paramiko.SSHException as e:
 126                 errors.append(e)
 127         raise SSHError("Invalid pkey: %s" % (errors))
 128
 129     def _get_client(self):
 130         if self._client:
 131             return self._client
 132         try:
 133             self._client = paramiko.SSHClient()
 134             self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 135             self._client.connect(self.host, username=self.user,
 136                                  port=self.port, pkey=self.pkey,
 137                                  key_filename=self.key_filename,
 138                                  password=self.password,
 139                                  allow_agent=False, look_for_keys=False,
 140                                  timeout=1)
 141             return self._client
 142         except Exception as e:
 143             message = ("Exception %(exception_type)s was raised "
 144                        "during connect. Exception value is: %(exception)r")
 145             self._client = False
 146             raise SSHError(message % {"exception": e,
 147                                       "exception_type": type(e)})
 148
 149     def close(self):
 150         self._client.close()
 151         self._client = False
 152
 153     def run(self, cmd, stdin=None, stdout=None, stderr=None,
 154             raise_on_error=True, timeout=3600,
 155             keep_stdin_open=False):
 156         """Execute specified command on the server.
 157
 158         :param cmd:             Command to be executed.
 159         :type cmd:              str
 160         :param stdin:           Open file or string to pass to stdin.
 161         :param stdout:          Open file to connect to stdout.
 162         :param stderr:          Open file to connect to stderr.
 163         :param raise_on_error:  If False then exit code will be return. If True
 164                                 then exception will be raized if non-zero code.
 165         :param timeout:         Timeout in seconds for command execution.
 166                                 Default 1 hour. No timeout if set to 0.
 167         :param keep_stdin_open: don't close stdin on empty reads
 168         :type keep_stdin_open:  bool
 169         """
 170
 171         client = self._get_client()
 172
 173         if isinstance(stdin, six.string_types):
 174             stdin = six.moves.StringIO(stdin)
 175
 176         return self._run(client, cmd, stdin=stdin, stdout=stdout,
 177                          stderr=stderr, raise_on_error=raise_on_error,
 178                          timeout=timeout,
 179                          keep_stdin_open=keep_stdin_open)
 180
 181     def _run(self, client, cmd, stdin=None, stdout=None, stderr=None,
 182              raise_on_error=True, timeout=3600,
 183              keep_stdin_open=False):
 184
 185         transport = client.get_transport()
 186         session = transport.open_session()
 187         session.exec_command(cmd)
 188         start_time = time.time()
 189
 190         data_to_send = ""
 191         stderr_data = None
 192
 193         # If we have data to be sent to stdin then `select' should also
 194         # check for stdin availability.
 195         if stdin and not stdin.closed:
 196             writes = [session]
 197         else:
 198             writes = []
 199
 200         while True:
 201             # Block until data can be read/write.
 202             r, w, e = select.select([session], writes, [session], 1)
 203
 204             if session.recv_ready():
 205                 data = session.recv(4096)
 206                 self.log.debug("stdout: %r", data)
 207                 if stdout is not None:
 208                     stdout.write(data)
 209                 continue
 210
 211             if session.recv_stderr_ready():
 212                 stderr_data = session.recv_stderr(4096)
 213                 self.log.debug("stderr: %r", stderr_data)
 214                 if stderr is not None:
 215                     stderr.write(stderr_data)
 216                 continue
 217
 218             if session.send_ready():
 219                 if stdin is not None and not stdin.closed:
 220                     if not data_to_send:
 221                         data_to_send = stdin.read(4096)
 222                         if not data_to_send:
 223                             # we may need to keep stdin open
 224                             if not keep_stdin_open:
 225                                 stdin.close()
 226                                 session.shutdown_write()
 227                                 writes = []
 228                     if data_to_send:
 229                         sent_bytes = session.send(data_to_send)
 230                         # LOG.debug("sent: %s" % data_to_send[:sent_bytes])
 231                         data_to_send = data_to_send[sent_bytes:]
 232
 233             if session.exit_status_ready():
 234                 break
 235
 236             if timeout and (time.time() - timeout) > start_time:
 237                 args = {"cmd": cmd, "host": self.host}
 238                 raise SSHTimeout("Timeout executing command "
 239                                  "'%(cmd)s' on host %(host)s" % args)
 240             if e:
 241                 raise SSHError("Socket error.")
 242
 243         exit_status = session.recv_exit_status()
 244         if 0 != exit_status and raise_on_error:
 245             fmt = "Command '%(cmd)s' failed with exit_status %(status)d."
 246             details = fmt % {"cmd": cmd, "status": exit_status}
 247             if stderr_data:
 248                 details += " Last stderr data: '%s'." % stderr_data
 249             raise SSHError(details)
 250         return exit_status
 251
 252     def execute(self, cmd, stdin=None, timeout=3600):
 253         """Execute the specified command on the server.
 254
 255         :param cmd:     Command to be executed.
 256         :param stdin:   Open file to be sent on process stdin.
 257         :param timeout: Timeout for execution of the command.
 258
 259         :returns: tuple (exit_status, stdout, stderr)
 260         """
 261         stdout = six.moves.StringIO()
 262         stderr = six.moves.StringIO()
 263
 264         exit_status = self.run(cmd, stderr=stderr,
 265                                stdout=stdout, stdin=stdin,
 266                                timeout=timeout, raise_on_error=False)
 267         stdout.seek(0)
 268         stderr.seek(0)
 269         return (exit_status, stdout.read(), stderr.read())
 270
 271     def wait(self, timeout=120, interval=1):
 272         """Wait for the host will be available via ssh."""
 273         start_time = time.time()
 274         while True:
 275             try:
 276                 return self.execute("uname")
 277             except (socket.error, SSHError) as e:
 278                 self.log.debug("Ssh is still unavailable: %r", e)
 279                 time.sleep(interval)
 280             if time.time() > (start_time + timeout):
 281                 raise SSHTimeout("Timeout waiting for '%s'", self.host)
 282
 283     def put(self, files, remote_path=b'.', recursive=False):
 284         client = self._get_client()
 285
 286         with SCPClient(client.get_transport()) as scp:
 287             scp.put(files, remote_path, recursive)
 288
 289     # keep shell running in the background, e.g. screen
 290     def send_command(self, command):
 291         client = self._get_client()
 292         client.exec_command(command, get_pty=True)