diff options
Diffstat (limited to 'yardstick/ssh.py')
| -rw-r--r-- | yardstick/ssh.py | 89 |
1 files changed, 71 insertions, 18 deletions
diff --git a/yardstick/ssh.py b/yardstick/ssh.py index d6ecffcb9..46d53b7d2 100644 --- a/yardstick/ssh.py +++ b/yardstick/ssh.py @@ -62,17 +62,16 @@ Eventlet: sshclient = eventlet.import_patched("yardstick.ssh") """ - +import os import select import socket import time +import logging import paramiko from scp import SCPClient import six -import logging -LOG = logging.getLogger(__name__) DEFAULT_PORT = 22 @@ -89,7 +88,7 @@ class SSH(object): """Represent ssh connection.""" def __init__(self, user, host, port=DEFAULT_PORT, pkey=None, - key_filename=None, password=None): + key_filename=None, password=None, name=None): """Initialize SSH client. :param user: ssh username @@ -99,6 +98,11 @@ 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 @@ -108,6 +112,13 @@ class SSH(object): 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) def _get_pkey(self, key): if isinstance(key, six.string_types): @@ -145,10 +156,12 @@ class SSH(object): 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): """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. @@ -156,6 +169,8 @@ 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 """ client = self._get_client() @@ -165,10 +180,12 @@ 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) 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): transport = client.get_transport() session = transport.open_session() @@ -191,14 +208,14 @@ class SSH(object): if session.recv_ready(): data = session.recv(4096) - LOG.debug("stdout: %r" % data) + 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) + self.log.debug("stderr: %r", stderr_data) if stderr is not None: stderr.write(stderr_data) continue @@ -208,13 +225,15 @@ class SSH(object): if not data_to_send: data_to_send = stdin.read(4096) 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 @@ -261,10 +280,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() @@ -276,3 +295,37 @@ 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) + + def _put_file_shell(self, localpath, remotepath, mode=None): + # quote to stop wordpslit + cmd = ['cat > "%s"' % remotepath] + if mode is not None: + # use -- so no options + cmd.append('chmod -- 0%o "%s"' % (mode, remotepath)) + + with open(localpath, "rb") as localfile: + # only chmod on successful cat + cmd = "&& ".join(cmd) + self.run(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 + """ + import socket + try: + self._put_file_sftp(localpath, remotepath, mode=mode) + except (paramiko.SSHException, socket.error): + self._put_file_shell(localpath, remotepath, mode=mode) |