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(io.RawIOBase):
  33         def write(chunk):
  34             if "error" in chunk:
  35                 email_admin(chunk)
  36
  37     ssh = SSH("root", "example.com")
  38     with PseudoFile() as p:
  39         ssh.run("tail -f /var/log/syslog", stdout=p, timeout=False)
  40
  41 Execute local script on remote side:
  42
  43     ssh = sshclient.SSH("user", "example.com")
  44
  45     with open("~/myscript.sh", "r") as stdin_file:
  46         status, out, err = ssh.execute('/bin/sh -s "arg1" "arg2"',
  47                                        stdin=stdin_file)
  48
  49 Upload file:
  50
  51     ssh = SSH("user", "example.com")
  52     # use rb for binary files
  53     with open("/store/file.gz", "rb") as stdin_file:
  54         ssh.run("cat > ~/upload/file.gz", stdin=stdin_file)
  55
  56 Eventlet:
  57
  58     eventlet.monkey_patch(select=True, time=True)
  59     or
  60     eventlet.monkey_patch()
  61     or
  62     sshclient = eventlet.import_patched("yardstick.ssh")
  63
  64 """
  65 from __future__ import absolute_import
  66 import os
  67 import select
  68 import socket
  69 import time
  70 import re
  71
  72 import logging
  73 import paramiko
  74 from oslo_utils import encodeutils
  75 from scp import SCPClient
  76 import six
  77
  78
  79 DEFAULT_PORT = 22
  80
  81
  82 class SSHError(Exception):
  83     pass
  84
  85
  86 class SSHTimeout(SSHError):
  87     pass
  88
  89
  90 class SSH(object):
  91     """Represent ssh connection."""
  92
  93     def __init__(self, user, host, port=DEFAULT_PORT, pkey=None,
  94                  key_filename=None, password=None, name=None):
  95         """Initialize SSH client.
  96
  97         :param user: ssh username
  98         :param host: hostname or ip address of remote ssh server
  99         :param port: remote ssh port
 100         :param pkey: RSA or DSS private key string or file object
 101         :param key_filename: private key filename
 102         :param password: password
 103         """
 104         self.name = name
 105         if name:
 106             self.log = logging.getLogger(__name__ + '.' + self.name)
 107         else:
 108             self.log = logging.getLogger(__name__)
 109
 110         self.user = user
 111         self.host = host
 112         # we may get text port from YAML, convert to int
 113         self.port = int(port)
 114         self.pkey = self._get_pkey(pkey) if pkey else None
 115         self.password = password
 116         self.key_filename = key_filename
 117         self._client = False
 118         # paramiko loglevel debug will output ssh protocl debug
 119         # we don't ever really want that unless we are debugging paramiko
 120         # ssh issues
 121         if os.environ.get("PARAMIKO_DEBUG", "").lower() == "true":
 122             logging.getLogger("paramiko").setLevel(logging.DEBUG)
 123         else:
 124             logging.getLogger("paramiko").setLevel(logging.WARN)
 125
 126     def _get_pkey(self, key):
 127         if isinstance(key, six.string_types):
 128             key = six.moves.StringIO(key)
 129         errors = []
 130         for key_class in (paramiko.rsakey.RSAKey, paramiko.dsskey.DSSKey):
 131             try:
 132                 return key_class.from_private_key(key)
 133             except paramiko.SSHException as e:
 134                 errors.append(e)
 135         raise SSHError("Invalid pkey: %s" % (errors))
 136
 137     def _get_client(self):
 138         if self._client:
 139             return self._client
 140         try:
 141             self._client = paramiko.SSHClient()
 142             self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 143             self._client.connect(self.host, username=self.user,
 144                                  port=self.port, pkey=self.pkey,
 145                                  key_filename=self.key_filename,
 146                                  password=self.password,
 147                                  allow_agent=False, look_for_keys=False,
 148                                  timeout=1)
 149             return self._client
 150         except Exception as e:
 151             message = ("Exception %(exception_type)s was raised "
 152                        "during connect. Exception value is: %(exception)r")
 153             self._client = False
 154             raise SSHError(message % {"exception": e,
 155                                       "exception_type": type(e)})
 156
 157     def close(self):
 158         self._client.close()
 159         self._client = False
 160
 161     def run(self, cmd, stdin=None, stdout=None, stderr=None,
 162             raise_on_error=True, timeout=3600,
 163             keep_stdin_open=False, pty=False):
 164         """Execute specified command on the server.
 165
 166         :param cmd:             Command to be executed.
 167         :type cmd:              str
 168         :param stdin:           Open file or string to pass to stdin.
 169         :param stdout:          Open file to connect to stdout.
 170         :param stderr:          Open file to connect to stderr.
 171         :param raise_on_error:  If False then exit code will be return. If True
 172                                 then exception will be raized if non-zero code.
 173         :param timeout:         Timeout in seconds for command execution.
 174                                 Default 1 hour. No timeout if set to 0.
 175         :param keep_stdin_open: don't close stdin on empty reads
 176         :type keep_stdin_open:  bool
 177         :param pty:             Request a pseudo terminal for this connection.
 178                                 This allows passing control characters.
 179                                 Default False.
 180         :type pty:              bool
 181         """
 182
 183         client = self._get_client()
 184
 185         if isinstance(stdin, six.string_types):
 186             stdin = six.moves.StringIO(stdin)
 187
 188         return self._run(client, cmd, stdin=stdin, stdout=stdout,
 189                          stderr=stderr, raise_on_error=raise_on_error,
 190                          timeout=timeout,
 191                          keep_stdin_open=keep_stdin_open, pty=pty)
 192
 193     def _run(self, client, cmd, stdin=None, stdout=None, stderr=None,
 194              raise_on_error=True, timeout=3600,
 195              keep_stdin_open=False, pty=False):
 196
 197         transport = client.get_transport()
 198         session = transport.open_session()
 199         if pty:
 200             session.get_pty()
 201         session.exec_command(cmd)
 202         start_time = time.time()
 203
 204         # encode on transmit, decode on receive
 205         data_to_send = encodeutils.safe_encode("", incoming='utf-8')
 206         stderr_data = None
 207
 208         # If we have data to be sent to stdin then `select' should also
 209         # check for stdin availability.
 210         if stdin and not stdin.closed:
 211             writes = [session]
 212         else:
 213             writes = []
 214
 215         while True:
 216             # Block until data can be read/write.
 217             r, w, e = select.select([session], writes, [session], 1)
 218
 219             if session.recv_ready():
 220                 data = encodeutils.safe_decode(session.recv(4096), 'utf-8')
 221                 self.log.debug("stdout: %r", data)
 222                 if stdout is not None:
 223                     stdout.write(data)
 224                 continue
 225
 226             if session.recv_stderr_ready():
 227                 stderr_data = encodeutils.safe_decode(
 228                     session.recv_stderr(4096), 'utf-8')
 229                 self.log.debug("stderr: %r", stderr_data)
 230                 if stderr is not None:
 231                     stderr.write(stderr_data)
 232                 continue
 233
 234             if session.send_ready():
 235                 if stdin is not None and not stdin.closed:
 236                     if not data_to_send:
 237                         stdin_txt = stdin.read(4096)
 238                         if stdin_txt is None:
 239                             stdin_txt = ''
 240                         data_to_send = encodeutils.safe_encode(
 241                             stdin_txt, incoming='utf-8')
 242                         if not data_to_send:
 243                             # we may need to keep stdin open
 244                             if not keep_stdin_open:
 245                                 stdin.close()
 246                                 session.shutdown_write()
 247                                 writes = []
 248                     if data_to_send:
 249                         sent_bytes = session.send(data_to_send)
 250                         # LOG.debug("sent: %s" % data_to_send[:sent_bytes])
 251                         data_to_send = data_to_send[sent_bytes:]
 252
 253             if session.exit_status_ready():
 254                 break
 255
 256             if timeout and (time.time() - timeout) > start_time:
 257                 args = {"cmd": cmd, "host": self.host}
 258                 raise SSHTimeout("Timeout executing command "
 259                                  "'%(cmd)s' on host %(host)s" % args)
 260             if e:
 261                 raise SSHError("Socket error.")
 262
 263         exit_status = session.recv_exit_status()
 264         if exit_status != 0 and raise_on_error:
 265             fmt = "Command '%(cmd)s' failed with exit_status %(status)d."
 266             details = fmt % {"cmd": cmd, "status": exit_status}
 267             if stderr_data:
 268                 details += " Last stderr data: '%s'." % stderr_data
 269             raise SSHError(details)
 270         return exit_status
 271
 272     def execute(self, cmd, stdin=None, timeout=3600):
 273         """Execute the specified command on the server.
 274
 275         :param cmd:     Command to be executed.
 276         :param stdin:   Open file to be sent on process stdin.
 277         :param timeout: Timeout for execution of the command.
 278
 279         :returns: tuple (exit_status, stdout, stderr)
 280         """
 281         stdout = six.moves.StringIO()
 282         stderr = six.moves.StringIO()
 283
 284         exit_status = self.run(cmd, stderr=stderr,
 285                                stdout=stdout, stdin=stdin,
 286                                timeout=timeout, raise_on_error=False)
 287         stdout.seek(0)
 288         stderr.seek(0)
 289         return (exit_status, stdout.read(), stderr.read())
 290
 291     def wait(self, timeout=120, interval=1):
 292         """Wait for the host will be available via ssh."""
 293         start_time = time.time()
 294         while True:
 295             try:
 296                 return self.execute("uname")
 297             except (socket.error, SSHError) as e:
 298                 self.log.debug("Ssh is still unavailable: %r", e)
 299                 time.sleep(interval)
 300             if time.time() > (start_time + timeout):
 301                 raise SSHTimeout("Timeout waiting for '%s'", self.host)
 302
 303     def put(self, files, remote_path=b'.', recursive=False):
 304         client = self._get_client()
 305
 306         with SCPClient(client.get_transport()) as scp:
 307             scp.put(files, remote_path, recursive)
 308
 309     # keep shell running in the background, e.g. screen
 310     def send_command(self, command):
 311         client = self._get_client()
 312         client.exec_command(command, get_pty=True)
 313
 314     def _put_file_sftp(self, localpath, remotepath, mode=None):
 315         client = self._get_client()
 316
 317         with client.open_sftp() as sftp:
 318             sftp.put(localpath, remotepath)
 319             if mode is None:
 320                 mode = 0o777 & os.stat(localpath).st_mode
 321             sftp.chmod(remotepath, mode)
 322
 323     TILDE_EXPANSIONS_RE = re.compile("(^~[^/]*/)?(.*)")
 324
 325     def _put_file_shell(self, localpath, remotepath, mode=None):
 326         # quote to stop wordpslit
 327         tilde, remotepath = self.TILDE_EXPANSIONS_RE.match(remotepath).groups()
 328         if not tilde:
 329             tilde = ''
 330         cmd = ['cat > %s"%s"' % (tilde, remotepath)]
 331         if mode is not None:
 332             # use -- so no options
 333             cmd.append('chmod -- 0%o %s"%s"' % (mode, tilde, remotepath))
 334
 335         with open(localpath, "rb") as localfile:
 336             # only chmod on successful cat
 337             self.run("&& ".join(cmd), stdin=localfile)
 338
 339     def put_file(self, localpath, remotepath, mode=None):
 340         """Copy specified local file to the server.
 341
 342         :param localpath:   Local filename.
 343         :param remotepath:  Remote filename.
 344         :param mode:        Permissions to set after upload
 345         """
 346         try:
 347             self._put_file_sftp(localpath, remotepath, mode=mode)
 348         except (paramiko.SSHException, socket.error):
 349             self._put_file_shell(localpath, remotepath, mode=mode)