X-Git-Url: https://gerrit.opnfv.org/gerrit/gitweb?a=blobdiff_plain;f=yardstick%2Fssh.py;h=a024cf64ae58ca35cf1506ec57530cf1c01d8f5c;hb=3d462aaa3d87a2298cf1f9a859e996d708f917fd;hp=2816a1c7db70a14227200201baee21b303b3c0fb;hpb=81d438404f189b1936777f1ec5abc34b6d171ffc;p=yardstick.git diff --git a/yardstick/ssh.py b/yardstick/ssh.py index 2816a1c7d..a024cf64a 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,20 +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 -import logging -LOG = logging.getLogger(__name__) +from yardstick.common.utils import try_int +from yardstick.network_services.utils import provision_tool class SSHError(Exception): @@ -81,8 +92,27 @@ class SSHTimeout(SSHError): class SSH(object): """Represent ssh connection.""" - def __init__(self, user, host, port=22, pkey=None, - key_filename=None, password=None): + SSH_PORT = paramiko.config.SSH_PORT + + @staticmethod + def gen_keys(key_filename, bit_count=2048): + rsa_key = paramiko.RSAKey.generate(bits=bit_count, progress_func=None) + rsa_key.write_private_key_file(key_filename) + print("Writing %s ..." % key_filename) + with open('.'.join([key_filename, "pub"]), "w") as pubkey_file: + pubkey_file.write(rsa_key.get_name()) + pubkey_file.write(' ') + pubkey_file.write(rsa_key.get_base64()) + pubkey_file.write('\n') + + @staticmethod + def get_class(): + # must return static class name, anything else refers to the calling class + # i.e. the subclass, not the superclass + return SSH + + def __init__(self, user, host, port=None, pkey=None, + key_filename=None, password=None, name=None): """Initialize SSH client. :param user: ssh username @@ -92,14 +122,51 @@ class SSH(object): :param key_filename: private key filename :param password: password """ + self.name = name + if name: + self.log = logging.getLogger(__name__ + '.' + self.name) + else: + self.log = logging.getLogger(__name__) self.user = user self.host = host - self.port = port + # 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 = try_int(port, self.SSH_PORT) self.pkey = self._get_pkey(pkey) if pkey else None self.password = password self.key_filename = key_filename self._client = False + # paramiko loglevel debug will output ssh protocl debug + # we don't ever really want that unless we are debugging paramiko + # ssh issues + if os.environ.get("PARAMIKO_DEBUG", "").lower() == "true": + logging.getLogger("paramiko").setLevel(logging.DEBUG) + else: + logging.getLogger("paramiko").setLevel(logging.WARN) + + @classmethod + def args_from_node(cls, node, overrides=None, defaults=None): + if overrides is None: + overrides = {} + if defaults is None: + defaults = {} + params = ChainMap(overrides, node, defaults) + return { + 'user': params['user'], + 'host': params['ip'], + 'port': params.get('ssh_port', cls.SSH_PORT), + 'pkey': params.get('pkey'), + 'key_filename': params.get('key_filename'), + 'password': params.get('password'), + 'name': params.get('name'), + } + + @classmethod + def from_node(cls, node, overrides=None, defaults=None): + return cls(**cls.args_from_node(node, overrides, defaults)) def _get_pkey(self, key): if isinstance(key, six.string_types): @@ -112,8 +179,12 @@ class SSH(object): errors.append(e) raise SSHError("Invalid pkey: %s" % (errors)) + @property + def is_connected(self): + return bool(self._client) + def _get_client(self): - if self._client: + if self.is_connected: return self._client try: self._client = paramiko.SSHClient() @@ -121,7 +192,9 @@ class SSH(object): self._client.connect(self.host, username=self.user, port=self.port, pkey=self.pkey, key_filename=self.key_filename, - password=self.password, timeout=1) + password=self.password, + allow_agent=False, look_for_keys=False, + timeout=1) return self._client except Exception as e: message = ("Exception %(exception_type)s was raised " @@ -130,15 +203,32 @@ class SSH(object): raise SSHError(message % {"exception": e, "exception_type": type(e)}) + def _make_dict(self): + return { + 'user': self.user, + 'host': self.host, + 'port': self.port, + 'pkey': self.pkey, + 'key_filename': self.key_filename, + 'password': self.password, + 'name': self.name, + } + + def copy(self): + return self.get_class()(**self._make_dict()) + def close(self): - self._client.close() - self._client = False + if self._client: + self._client.close() + self._client = False def run(self, cmd, stdin=None, stdout=None, stderr=None, - raise_on_error=True, timeout=3600): + raise_on_error=True, timeout=3600, + keep_stdin_open=False, pty=False): """Execute specified command on the server. :param cmd: Command to be executed. + :type cmd: str :param stdin: Open file or string to pass to stdin. :param stdout: Open file to connect to stdout. :param stderr: Open file to connect to stderr. @@ -146,6 +236,12 @@ class SSH(object): then exception will be raized if non-zero code. :param timeout: Timeout in seconds for command execution. 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() @@ -155,17 +251,22 @@ class SSH(object): return self._run(client, cmd, stdin=stdin, stdout=stdout, stderr=stderr, raise_on_error=raise_on_error, - timeout=timeout) + timeout=timeout, + 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): + raise_on_error=True, timeout=3600, + 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 @@ -180,15 +281,16 @@ class SSH(object): r, w, e = select.select([session], writes, [session], 1) if session.recv_ready(): - data = session.recv(4096) - LOG.debug("stdout: %r" % data) + 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) - LOG.debug("stderr: %r" % stderr_data) + 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) continue @@ -196,15 +298,21 @@ 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: - stdin.close() - session.shutdown_write() - writes = [] - continue - sent_bytes = session.send(data_to_send) - # LOG.debug("sent: %s" % data_to_send[:sent_bytes]) - data_to_send = data_to_send[sent_bytes:] + # we may need to keep stdin open + if not keep_stdin_open: + stdin.close() + session.shutdown_write() + writes = [] + if data_to_send: + sent_bytes = session.send(data_to_send) + # LOG.debug("sent: %s" % data_to_send[:sent_bytes]) + data_to_send = data_to_send[sent_bytes:] if session.exit_status_ready(): break @@ -217,7 +325,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: @@ -242,7 +350,7 @@ class SSH(object): timeout=timeout, raise_on_error=False) stdout.seek(0) stderr.seek(0) - return (exit_status, stdout.read(), stderr.read()) + return exit_status, stdout.read(), stderr.read() def wait(self, timeout=120, interval=1): """Wait for the host will be available via ssh.""" @@ -251,10 +359,10 @@ class SSH(object): try: return self.execute("uname") except (socket.error, SSHError) as e: - LOG.debug("Ssh is still unavailable: %r" % e) + self.log.debug("Ssh is still unavailable: %r", e) time.sleep(interval) if time.time() > (start_time + timeout): - raise SSHTimeout("Timeout waiting for '%s'" % self.host) + raise SSHTimeout("Timeout waiting for '%s'", self.host) def put(self, files, remote_path=b'.', recursive=False): client = self._get_client() @@ -266,3 +374,119 @@ 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) + + def provision_tool(self, tool_path, tool_file=None): + return provision_tool(self, tool_path, tool_file) + + def put_file_obj(self, file_obj, remotepath, mode=None): + client = self._get_client() + + with client.open_sftp() as sftp: + sftp.putfo(file_obj, remotepath) + if mode is not None: + sftp.chmod(remotepath, mode) + + def get_file_obj(self, remotepath, file_obj): + client = self._get_client() + + with client.open_sftp() as sftp: + sftp.getfo(remotepath, file_obj) + + +class AutoConnectSSH(SSH): + + def __init__(self, user, host, port=None, pkey=None, + key_filename=None, password=None, name=None, wait=False): + super(AutoConnectSSH, self).__init__(user, host, port, pkey, key_filename, password, name) + self._wait = wait + + def _make_dict(self): + data = super(AutoConnectSSH, self)._make_dict() + data.update({ + 'wait': self._wait + }) + return data + + def _connect(self): + if not self.is_connected: + self._get_client() + if self._wait: + self.wait() + + def drop_connection(self): + """ Don't close anything, just force creation of a new client """ + self._client = False + + def execute(self, cmd, stdin=None, timeout=3600): + self._connect() + return super(AutoConnectSSH, self).execute(cmd, stdin, timeout) + + def run(self, cmd, stdin=None, stdout=None, stderr=None, + raise_on_error=True, timeout=3600, + keep_stdin_open=False, pty=False): + self._connect() + return super(AutoConnectSSH, self).run(cmd, stdin, stdout, stderr, raise_on_error, + timeout, keep_stdin_open, pty) + + def put(self, files, remote_path=b'.', recursive=False): + self._connect() + return super(AutoConnectSSH, self).put(files, remote_path, recursive) + + def put_file(self, local_path, remote_path, mode=None): + self._connect() + return super(AutoConnectSSH, self).put_file(local_path, remote_path, mode) + + def put_file_obj(self, file_obj, remote_path, mode=None): + self._connect() + return super(AutoConnectSSH, self).put_file_obj(file_obj, remote_path, mode) + + def get_file_obj(self, remote_path, file_obj): + self._connect() + return super(AutoConnectSSH, self).get_file_obj(remote_path, file_obj) + + def provision_tool(self, tool_path, tool_file=None): + self._connect() + return super(AutoConnectSSH, self).provision_tool(tool_path, tool_file) + + @staticmethod + def get_class(): + # must return static class name, anything else refers to the calling class + # i.e. the subclass, not the superclass + return AutoConnectSSH