1 # Copyright 2013: Mirantis Inc.
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
8 # http://www.apache.org/licenses/LICENSE-2.0
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
16 # yardstick comment: this is a modified copy of rally/rally/common/sshutils.py
18 """High level ssh library.
22 Execute command and get output:
24 ssh = sshclient.SSH("root", "example.com", port=33)
25 status, stdout, stderr = ssh.execute("ps ax")
27 raise Exception("Command failed with non-zero status.")
28 print(stdout.splitlines())
30 Execute command with huge output:
32 class PseudoFile(io.RawIOBase):
37 ssh = SSH("root", "example.com")
38 with PseudoFile() as p:
39 ssh.run("tail -f /var/log/syslog", stdout=p, timeout=False)
41 Execute local script on remote side:
43 ssh = sshclient.SSH("user", "example.com")
45 with open("~/myscript.sh", "r") as stdin_file:
46 status, out, err = ssh.execute('/bin/sh -s "arg1" "arg2"',
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)
58 eventlet.monkey_patch(select=True, time=True)
60 eventlet.monkey_patch()
62 sshclient = eventlet.import_patched("yardstick.ssh")
74 from chainmap import ChainMap
75 from oslo_utils import encodeutils
76 from scp import SCPClient
79 from yardstick.common import exceptions
80 from yardstick.common.utils import try_int, NON_NONE_DEFAULT, make_dict_from_map
81 from yardstick.network_services.utils import provision_tool
84 def convert_key_to_str(key):
85 if not isinstance(key, (paramiko.RSAKey, paramiko.DSSKey)):
88 key.write_private_key(k)
93 """Represent ssh connection."""
95 SSH_PORT = paramiko.config.SSH_PORT
96 DEFAULT_WAIT_TIMEOUT = 120
99 def gen_keys(key_filename, bit_count=2048):
100 rsa_key = paramiko.RSAKey.generate(bits=bit_count, progress_func=None)
101 rsa_key.write_private_key_file(key_filename)
102 print("Writing %s ..." % key_filename)
103 with open('.'.join([key_filename, "pub"]), "w") as pubkey_file:
104 pubkey_file.write(rsa_key.get_name())
105 pubkey_file.write(' ')
106 pubkey_file.write(rsa_key.get_base64())
107 pubkey_file.write('\n')
111 # must return static class name, anything else refers to the calling class
112 # i.e. the subclass, not the superclass
116 def get_arg_key_map(cls):
118 'user': ('user', NON_NONE_DEFAULT),
119 'host': ('ip', NON_NONE_DEFAULT),
120 'port': ('ssh_port', cls.SSH_PORT),
121 'pkey': ('pkey', None),
122 'key_filename': ('key_filename', None),
123 'password': ('password', None),
124 'name': ('name', None),
127 def __init__(self, user, host, port=None, pkey=None,
128 key_filename=None, password=None, name=None):
129 """Initialize SSH client.
131 :param user: ssh username
132 :param host: hostname or ip address of remote ssh server
133 :param port: remote ssh port
134 :param pkey: RSA or DSS private key string or file object
135 :param key_filename: private key filename
136 :param password: password
140 self.log = logging.getLogger(__name__ + '.' + self.name)
142 self.log = logging.getLogger(__name__)
144 self.wait_timeout = self.DEFAULT_WAIT_TIMEOUT
147 # everybody wants to debug this in the caller, do it here instead
148 self.log.debug("user:%s host:%s", user, host)
150 # we may get text port from YAML, convert to int
151 self.port = try_int(port, self.SSH_PORT)
152 self.pkey = self._get_pkey(pkey) if pkey else None
153 self.password = password
154 self.key_filename = key_filename
156 # paramiko loglevel debug will output ssh protocl debug
157 # we don't ever really want that unless we are debugging paramiko
159 if os.environ.get("PARAMIKO_DEBUG", "").lower() == "true":
160 logging.getLogger("paramiko").setLevel(logging.DEBUG)
162 logging.getLogger("paramiko").setLevel(logging.WARN)
165 def args_from_node(cls, node, overrides=None, defaults=None):
166 if overrides is None:
171 params = ChainMap(overrides, node, defaults)
172 return make_dict_from_map(params, cls.get_arg_key_map())
175 def from_node(cls, node, overrides=None, defaults=None):
176 return cls(**cls.args_from_node(node, overrides, defaults))
178 def _get_pkey(self, key):
179 if isinstance(key, six.string_types):
180 key = six.moves.StringIO(key)
182 for key_class in (paramiko.rsakey.RSAKey, paramiko.dsskey.DSSKey):
184 return key_class.from_private_key(key)
185 except paramiko.SSHException as e:
187 raise exceptions.SSHError(error_msg='Invalid pkey: %s' % errors)
190 def is_connected(self):
191 return bool(self._client)
193 def _get_client(self):
194 if self.is_connected:
197 self._client = paramiko.SSHClient()
198 self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
199 self._client.connect(self.host, username=self.user,
200 port=self.port, pkey=self.pkey,
201 key_filename=self.key_filename,
202 password=self.password,
203 allow_agent=False, look_for_keys=False,
206 except Exception as e:
207 message = ("Exception %(exception_type)s was raised "
208 "during connect. Exception value is: %(exception)r" %
209 {"exception": e, "exception_type": type(e)})
211 raise exceptions.SSHError(error_msg=message)
213 def _make_dict(self):
219 'key_filename': self.key_filename,
220 'password': self.password,
225 return self.get_class()(**self._make_dict())
232 def run(self, cmd, stdin=None, stdout=None, stderr=None,
233 raise_on_error=True, timeout=3600,
234 keep_stdin_open=False, pty=False):
235 """Execute specified command on the server.
237 :param cmd: Command to be executed.
239 :param stdin: Open file or string to pass to stdin.
240 :param stdout: Open file to connect to stdout.
241 :param stderr: Open file to connect to stderr.
242 :param raise_on_error: If False then exit code will be return. If True
243 then exception will be raized if non-zero code.
244 :param timeout: Timeout in seconds for command execution.
245 Default 1 hour. No timeout if set to 0.
246 :param keep_stdin_open: don't close stdin on empty reads
247 :type keep_stdin_open: bool
248 :param pty: Request a pseudo terminal for this connection.
249 This allows passing control characters.
254 client = self._get_client()
256 if isinstance(stdin, six.string_types):
257 stdin = six.moves.StringIO(stdin)
259 return self._run(client, cmd, stdin=stdin, stdout=stdout,
260 stderr=stderr, raise_on_error=raise_on_error,
262 keep_stdin_open=keep_stdin_open, pty=pty)
264 def _run(self, client, cmd, stdin=None, stdout=None, stderr=None,
265 raise_on_error=True, timeout=3600,
266 keep_stdin_open=False, pty=False):
268 transport = client.get_transport()
269 session = transport.open_session()
272 session.exec_command(cmd)
273 start_time = time.time()
275 # encode on transmit, decode on receive
276 data_to_send = encodeutils.safe_encode("", incoming='utf-8')
279 # If we have data to be sent to stdin then `select' should also
280 # check for stdin availability.
281 if stdin and not stdin.closed:
287 # Block until data can be read/write.
288 e = select.select([session], writes, [session], 1)[2]
290 if session.recv_ready():
291 data = encodeutils.safe_decode(session.recv(4096), 'utf-8')
292 self.log.debug("stdout: %r", data)
293 if stdout is not None:
297 if session.recv_stderr_ready():
298 stderr_data = encodeutils.safe_decode(
299 session.recv_stderr(4096), 'utf-8')
300 self.log.debug("stderr: %r", stderr_data)
301 if stderr is not None:
302 stderr.write(stderr_data)
305 if session.send_ready():
306 if stdin is not None and not stdin.closed:
308 stdin_txt = stdin.read(4096)
309 if stdin_txt is None:
311 data_to_send = encodeutils.safe_encode(
312 stdin_txt, incoming='utf-8')
314 # we may need to keep stdin open
315 if not keep_stdin_open:
317 session.shutdown_write()
320 sent_bytes = session.send(data_to_send)
321 # LOG.debug("sent: %s" % data_to_send[:sent_bytes])
322 data_to_send = data_to_send[sent_bytes:]
324 if session.exit_status_ready():
327 if timeout and (time.time() - timeout) > start_time:
328 message = ('Timeout executing command %(cmd)s on host %(host)s'
329 % {"cmd": cmd, "host": self.host})
330 raise exceptions.SSHTimeout(error_msg=message)
332 raise exceptions.SSHError(error_msg='Socket error')
334 exit_status = session.recv_exit_status()
335 if exit_status != 0 and raise_on_error:
336 fmt = "Command '%(cmd)s' failed with exit_status %(status)d."
337 details = fmt % {"cmd": cmd, "status": exit_status}
339 details += " Last stderr data: '%s'." % stderr_data
340 raise exceptions.SSHError(error_msg=details)
343 def execute(self, cmd, stdin=None, timeout=3600, raise_on_error=False):
344 """Execute the specified command on the server.
346 :param cmd: (str) Command to be executed.
347 :param stdin: (StringIO) Open file to be sent on process stdin.
348 :param timeout: (int) Timeout for execution of the command.
349 :param raise_on_error: (bool) If True, then an SSHError will be raised
350 when non-zero exit code.
352 :returns: tuple (exit_status, stdout, stderr)
354 stdout = six.moves.StringIO()
355 stderr = six.moves.StringIO()
357 exit_status = self.run(cmd, stderr=stderr,
358 stdout=stdout, stdin=stdin,
359 timeout=timeout, raise_on_error=raise_on_error)
362 return exit_status, stdout.read(), stderr.read()
364 def wait(self, timeout=None, interval=1):
365 """Wait for the host will be available via ssh."""
367 timeout = self.wait_timeout
369 end_time = time.time() + timeout
372 return self.execute("uname")
373 except (socket.error, exceptions.SSHError) as e:
374 self.log.debug("Ssh is still unavailable: %r", e)
376 if time.time() > end_time:
377 raise exceptions.SSHTimeout(
378 error_msg='Timeout waiting for "%s"' % self.host)
380 def put(self, files, remote_path=b'.', recursive=False):
381 client = self._get_client()
383 with SCPClient(client.get_transport()) as scp:
384 scp.put(files, remote_path, recursive)
386 def get(self, remote_path, local_path='/tmp/', recursive=True):
387 client = self._get_client()
389 with SCPClient(client.get_transport()) as scp:
390 scp.get(remote_path, local_path, recursive)
392 # keep shell running in the background, e.g. screen
393 def send_command(self, command):
394 client = self._get_client()
395 client.exec_command(command, get_pty=True)
397 def _put_file_sftp(self, localpath, remotepath, mode=None):
398 client = self._get_client()
400 with client.open_sftp() as sftp:
401 sftp.put(localpath, remotepath)
403 mode = 0o777 & os.stat(localpath).st_mode
404 sftp.chmod(remotepath, mode)
406 TILDE_EXPANSIONS_RE = re.compile("(^~[^/]*/)?(.*)")
408 def _put_file_shell(self, localpath, remotepath, mode=None):
409 # quote to stop wordpslit
410 tilde, remotepath = self.TILDE_EXPANSIONS_RE.match(remotepath).groups()
413 cmd = ['cat > %s"%s"' % (tilde, remotepath)]
415 # use -- so no options
416 cmd.append('chmod -- 0%o %s"%s"' % (mode, tilde, remotepath))
418 with open(localpath, "rb") as localfile:
419 # only chmod on successful cat
420 self.run("&& ".join(cmd), stdin=localfile)
422 def put_file(self, localpath, remotepath, mode=None):
423 """Copy specified local file to the server.
425 :param localpath: Local filename.
426 :param remotepath: Remote filename.
427 :param mode: Permissions to set after upload
430 self._put_file_sftp(localpath, remotepath, mode=mode)
431 except (paramiko.SSHException, socket.error):
432 self._put_file_shell(localpath, remotepath, mode=mode)
434 def provision_tool(self, tool_path, tool_file=None):
435 return provision_tool(self, tool_path, tool_file)
437 def put_file_obj(self, file_obj, remotepath, mode=None):
438 client = self._get_client()
440 with client.open_sftp() as sftp:
441 sftp.putfo(file_obj, remotepath)
443 sftp.chmod(remotepath, mode)
445 def get_file_obj(self, remotepath, file_obj):
446 client = self._get_client()
448 with client.open_sftp() as sftp:
449 sftp.getfo(remotepath, file_obj)
452 class AutoConnectSSH(SSH):
455 def get_arg_key_map(cls):
456 arg_key_map = super(AutoConnectSSH, cls).get_arg_key_map()
457 arg_key_map['wait'] = ('wait', True)
460 # always wait or we will get OpenStack SSH errors
461 def __init__(self, user, host, port=None, pkey=None,
462 key_filename=None, password=None, name=None, wait=True):
463 super(AutoConnectSSH, self).__init__(user, host, port, pkey, key_filename, password, name)
464 if wait and wait is not True:
465 self.wait_timeout = int(wait)
467 def _make_dict(self):
468 data = super(AutoConnectSSH, self)._make_dict()
470 'wait': self.wait_timeout
475 if not self.is_connected:
477 timeout = self.wait_timeout
479 end_time = time.time() + timeout
482 return self._get_client()
483 except (socket.error, exceptions.SSHError) as e:
484 self.log.debug("Ssh is still unavailable: %r", e)
486 if time.time() > end_time:
487 raise exceptions.SSHTimeout(
488 error_msg='Timeout waiting for "%s"' % self.host)
490 def drop_connection(self):
491 """ Don't close anything, just force creation of a new client """
494 def execute(self, cmd, stdin=None, timeout=3600, raise_on_error=False):
496 return super(AutoConnectSSH, self).execute(cmd, stdin, timeout,
499 def run(self, cmd, stdin=None, stdout=None, stderr=None,
500 raise_on_error=True, timeout=3600,
501 keep_stdin_open=False, pty=False):
503 return super(AutoConnectSSH, self).run(cmd, stdin, stdout, stderr, raise_on_error,
504 timeout, keep_stdin_open, pty)
506 def put(self, files, remote_path=b'.', recursive=False):
508 return super(AutoConnectSSH, self).put(files, remote_path, recursive)
510 def put_file(self, local_path, remote_path, mode=None):
512 return super(AutoConnectSSH, self).put_file(local_path, remote_path, mode)
514 def put_file_obj(self, file_obj, remote_path, mode=None):
516 return super(AutoConnectSSH, self).put_file_obj(file_obj, remote_path, mode)
518 def get_file_obj(self, remote_path, file_obj):
520 return super(AutoConnectSSH, self).get_file_obj(remote_path, file_obj)
522 def provision_tool(self, tool_path, tool_file=None):
524 return super(AutoConnectSSH, self).provision_tool(tool_path, tool_file)
528 # must return static class name, anything else refers to the calling class
529 # i.e. the subclass, not the superclass
530 return AutoConnectSSH