Merge "Bugfix: network_utilization: parse error"
[yardstick.git] / 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)