X-Git-Url: https://gerrit.opnfv.org/gerrit/gitweb?a=blobdiff_plain;f=yardstick%2Fssh.py;h=cf9adf0dc6d5c994fbc57301b18e94a3fd459885;hb=33b439426dac57bc8448df1bb485138550235740;hp=8485dccd098ccf08ae82d3722efe5c35756368f4;hpb=fad4811d1a722f2577f158840b7c8e61489e2e90;p=yardstick.git diff --git a/yardstick/ssh.py b/yardstick/ssh.py index 8485dccd0..cf9adf0dc 100644 --- a/yardstick/ssh.py +++ b/yardstick/ssh.py @@ -25,28 +25,33 @@ Execute command and get output: status, stdout, stderr = ssh.execute("ps ax") if status: raise Exception("Command failed with non-zero status.") - print stdout.splitlines() + print(stdout.splitlines()) Execute command with huge output: - class PseudoFile(object): + class PseudoFile(io.RawIOBase): def write(chunk): if "error" in chunk: email_admin(chunk) - ssh = sshclient.SSH("root", "example.com") - ssh.run("tail -f /var/log/syslog", stdout=PseudoFile(), timeout=False) + ssh = SSH("root", "example.com") + with PseudoFile() as p: + ssh.run("tail -f /var/log/syslog", stdout=p, timeout=False) Execute local script on remote side: ssh = sshclient.SSH("user", "example.com") - status, out, err = ssh.execute("/bin/sh -s arg1 arg2", - stdin=open("~/myscript.sh", "r")) + + with open("~/myscript.sh", "r") as stdin_file: + status, out, err = ssh.execute('/bin/sh -s "arg1" "arg2"', + stdin=stdin_file) Upload file: - ssh = sshclient.SSH("user", "example.com") - ssh.run("cat > ~/upload/file.gz", stdin=open("/store/file.gz", "rb")) + ssh = SSH("user", "example.com") + # use rb for binary files + with open("/store/file.gz", "rb") as stdin_file: + ssh.run("cat > ~/upload/file.gz", stdin=stdin_file) Eventlet: @@ -54,21 +59,26 @@ Eventlet: or eventlet.monkey_patch() or - sshclient = eventlet.import_patched("opentstack.common.sshclient") + sshclient = eventlet.import_patched("yardstick.ssh") """ +from __future__ import absolute_import import os import select import socket import time +import re import logging + import paramiko +from chainmap import ChainMap +from oslo_utils import encodeutils from scp import SCPClient import six -DEFAULT_PORT = 22 +SSH_PORT = paramiko.config.SSH_PORT class SSHError(Exception): @@ -82,7 +92,7 @@ class SSHTimeout(SSHError): class SSH(object): """Represent ssh connection.""" - def __init__(self, user, host, port=DEFAULT_PORT, pkey=None, + def __init__(self, user, host, port=SSH_PORT, pkey=None, key_filename=None, password=None, name=None): """Initialize SSH client. @@ -101,6 +111,9 @@ class SSH(object): self.user = user self.host = host + # everybody wants to debug this in the caller, do it here instead + self.log.debug("user:%s host:%s", user, host) + # we may get text port from YAML, convert to int self.port = int(port) self.pkey = self._get_pkey(pkey) if pkey else None @@ -115,6 +128,23 @@ class SSH(object): else: logging.getLogger("paramiko").setLevel(logging.WARN) + @classmethod + def from_node(cls, node, overrides=None, defaults=None): + if overrides is None: + overrides = {} + if defaults is None: + defaults = {} + params = ChainMap(overrides, node, defaults) + return cls( + user=params['user'], + host=params['ip'], + # paramiko doesn't like None default, requires SSH_PORT default + port=params.get('ssh_port', SSH_PORT), + pkey=params.get('pkey'), + key_filename=params.get('key_filename'), + password=params.get('password'), + name=params.get('name')) + def _get_pkey(self, key): if isinstance(key, six.string_types): key = six.moves.StringIO(key) @@ -152,7 +182,7 @@ class SSH(object): def run(self, cmd, stdin=None, stdout=None, stderr=None, raise_on_error=True, timeout=3600, - keep_stdin_open=False): + keep_stdin_open=False, pty=False): """Execute specified command on the server. :param cmd: Command to be executed. @@ -166,6 +196,10 @@ class SSH(object): Default 1 hour. No timeout if set to 0. :param keep_stdin_open: don't close stdin on empty reads :type keep_stdin_open: bool + :param pty: Request a pseudo terminal for this connection. + This allows passing control characters. + Default False. + :type pty: bool """ client = self._get_client() @@ -176,18 +210,21 @@ class SSH(object): return self._run(client, cmd, stdin=stdin, stdout=stdout, stderr=stderr, raise_on_error=raise_on_error, timeout=timeout, - keep_stdin_open=keep_stdin_open) + keep_stdin_open=keep_stdin_open, pty=pty) def _run(self, client, cmd, stdin=None, stdout=None, stderr=None, raise_on_error=True, timeout=3600, - keep_stdin_open=False): + keep_stdin_open=False, pty=False): transport = client.get_transport() session = transport.open_session() + if pty: + session.get_pty() session.exec_command(cmd) start_time = time.time() - data_to_send = "" + # encode on transmit, decode on receive + data_to_send = encodeutils.safe_encode("", incoming='utf-8') stderr_data = None # If we have data to be sent to stdin then `select' should also @@ -202,14 +239,15 @@ class SSH(object): r, w, e = select.select([session], writes, [session], 1) if session.recv_ready(): - data = session.recv(4096) + data = encodeutils.safe_decode(session.recv(4096), 'utf-8') self.log.debug("stdout: %r", data) if stdout is not None: stdout.write(data) continue if session.recv_stderr_ready(): - stderr_data = session.recv_stderr(4096) + stderr_data = encodeutils.safe_decode( + session.recv_stderr(4096), 'utf-8') self.log.debug("stderr: %r", stderr_data) if stderr is not None: stderr.write(stderr_data) @@ -218,7 +256,11 @@ class SSH(object): if session.send_ready(): if stdin is not None and not stdin.closed: if not data_to_send: - data_to_send = stdin.read(4096) + stdin_txt = stdin.read(4096) + if stdin_txt is None: + stdin_txt = '' + data_to_send = encodeutils.safe_encode( + stdin_txt, incoming='utf-8') if not data_to_send: # we may need to keep stdin open if not keep_stdin_open: @@ -241,7 +283,7 @@ class SSH(object): raise SSHError("Socket error.") exit_status = session.recv_exit_status() - if 0 != exit_status and raise_on_error: + if exit_status != 0 and raise_on_error: fmt = "Command '%(cmd)s' failed with exit_status %(status)d." details = fmt % {"cmd": cmd, "status": exit_status} if stderr_data: @@ -290,3 +332,40 @@ class SSH(object): def send_command(self, command): client = self._get_client() client.exec_command(command, get_pty=True) + + def _put_file_sftp(self, localpath, remotepath, mode=None): + client = self._get_client() + + with client.open_sftp() as sftp: + sftp.put(localpath, remotepath) + if mode is None: + mode = 0o777 & os.stat(localpath).st_mode + sftp.chmod(remotepath, mode) + + TILDE_EXPANSIONS_RE = re.compile("(^~[^/]*/)?(.*)") + + def _put_file_shell(self, localpath, remotepath, mode=None): + # quote to stop wordpslit + tilde, remotepath = self.TILDE_EXPANSIONS_RE.match(remotepath).groups() + if not tilde: + tilde = '' + cmd = ['cat > %s"%s"' % (tilde, remotepath)] + if mode is not None: + # use -- so no options + cmd.append('chmod -- 0%o %s"%s"' % (mode, tilde, remotepath)) + + with open(localpath, "rb") as localfile: + # only chmod on successful cat + self.run("&& ".join(cmd), stdin=localfile) + + def put_file(self, localpath, remotepath, mode=None): + """Copy specified local file to the server. + + :param localpath: Local filename. + :param remotepath: Remote filename. + :param mode: Permissions to set after upload + """ + try: + self._put_file_sftp(localpath, remotepath, mode=mode) + except (paramiko.SSHException, socket.error): + self._put_file_shell(localpath, remotepath, mode=mode)