• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2019 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Ssh Utilities."""
15from __future__ import print_function
16import logging
17
18import re
19import subprocess
20import sys
21import threading
22
23from acloud import errors
24from acloud.internal import constants
25from acloud.internal.lib import utils
26
27logger = logging.getLogger(__name__)
28
29_SSH_CMD = ("-i %(rsa_key_file)s -o LogLevel=ERROR "
30            "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no")
31_SSH_IDENTITY = "-l %(login_user)s %(ip_addr)s"
32_SSH_CMD_MAX_RETRY = 5
33_SSH_CMD_RETRY_SLEEP = 3
34_CONNECTION_TIMEOUT = 10
35_MAX_REPORTED_ERROR_LINES = 10
36_ERROR_MSG_RE = re.compile(r".*]\s*\"(?:message|response)\"\s:\s\"(?P<content>.*)\"")
37_ERROR_MSG_TO_QUOTE_RE = r"(\\u2019)|(\\u2018)"
38_ERROR_MSG_DEL_STYLE_RE = r"(<style.+\/style>)"
39_ERROR_MSG_DEL_TAGS_RE = (r"(<[\/]*(a|b|p|span|ins|code|title)>)|"
40                          r"(<(a|span|meta|html|!)[^>]*>)")
41
42
43def _SshCallWait(cmd, timeout=None):
44    """Runs a single SSH command.
45
46    - SSH returns code 0 for "Successful execution".
47    - Use wait() until the process is complete without receiving any output.
48
49    Args:
50        cmd: String of the full SSH command to run, including the SSH binary
51             and its arguments.
52        timeout: Optional integer, number of seconds to give
53
54    Returns:
55        An exit status of 0 indicates that it ran successfully.
56    """
57    logger.info("Running command \"%s\"", cmd)
58    process = subprocess.Popen(cmd, shell=True, stdin=None,
59                               stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
60    if timeout:
61        # TODO: if process is killed, out error message to log.
62        timer = threading.Timer(timeout, process.kill)
63        timer.start()
64    process.wait()
65    if timeout:
66        timer.cancel()
67    return process.returncode
68
69
70def _SshCall(cmd, timeout=None):
71    """Runs a single SSH command.
72
73    - SSH returns code 0 for "Successful execution".
74    - Use communicate() until the process and the child thread are complete.
75
76    Args:
77        cmd: String of the full SSH command to run, including the SSH binary
78             and its arguments.
79        timeout: Optional integer, number of seconds to give
80
81    Returns:
82        An exit status of 0 indicates that it ran successfully.
83    """
84    logger.info("Running command \"%s\"", cmd)
85    process = subprocess.Popen(cmd, shell=True, stdin=None,
86                               stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
87    if timeout:
88        # TODO: if process is killed, out error message to log.
89        timer = threading.Timer(timeout, process.kill)
90        timer.start()
91    process.communicate()
92    if timeout:
93        timer.cancel()
94    return process.returncode
95
96
97def _SshLogOutput(cmd, timeout=None, show_output=False):
98    """Runs a single SSH command while logging its output and processes its return code.
99
100    Output is streamed to the log at the debug level for more interactive debugging.
101    SSH returns error code 255 for "failed to connect", so this is interpreted as a failure in
102    SSH rather than a failure on the target device and this is converted to a different exception
103    type.
104
105    Args:
106        cmd: String of the full SSH command to run, including the SSH binary and its arguments.
107        timeout: Optional integer, number of seconds to give.
108        show_output: Boolean, True to show command output in screen.
109
110    Raises:
111        errors.DeviceConnectionError: Failed to connect to the GCE instance.
112        subprocess.CalledProcessError: The process exited with an error on the instance.
113        errors.LaunchCVDFail: Happened on launch_cvd with specific pattern of error message.
114    """
115    # Use "exec" to let cmd to inherit the shell process, instead of having the
116    # shell launch a child process which does not get killed.
117    cmd = "exec " + cmd
118    logger.info("Running command \"%s\"", cmd)
119    process = subprocess.Popen(cmd, shell=True, stdin=None,
120                               stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
121                               universal_newlines=True)
122    if timeout:
123        # TODO: if process is killed, out error message to log.
124        timer = threading.Timer(timeout, process.kill)
125        timer.start()
126    stdout, _ = process.communicate()
127    if stdout:
128        if show_output or process.returncode != 0:
129            print(stdout.strip(), file=sys.stderr)
130        else:
131            # fetch_cvd and launch_cvd can be noisy, so left at debug
132            logger.debug(stdout.strip())
133    if timeout:
134        timer.cancel()
135    if process.returncode == 255:
136        raise errors.DeviceConnectionError(
137            "Failed to send command to instance (%(ssh_cmd)s)\n"
138            "Error message: %(error_message)s" % {
139                "ssh_cmd": cmd,
140                "error_message": _GetErrorMessage(stdout)}
141        )
142    if process.returncode != 0:
143        if constants.ERROR_MSG_VNC_NOT_SUPPORT in stdout:
144            raise errors.LaunchCVDFail(constants.ERROR_MSG_VNC_NOT_SUPPORT)
145        if constants.ERROR_MSG_WEBRTC_NOT_SUPPORT in stdout:
146            raise errors.LaunchCVDFail(constants.ERROR_MSG_WEBRTC_NOT_SUPPORT)
147        raise subprocess.CalledProcessError(process.returncode, cmd)
148
149
150def _GetErrorMessage(stdout):
151    """Get error message.
152
153    Fetch the content of "message" or "response" from the ssh output and filter
154    unused content then log into report. Once the two fields didn't match, to
155    log last _MAX_REPORTED_ERROR_LINES lines into report.
156
157    Args:
158        stdout: String of the ssh output.
159
160    Returns:
161        String of the formatted ssh output.
162    """
163    matches = _ERROR_MSG_RE.finditer(stdout)
164    for match in matches:
165        return _FilterUnusedContent(match.group("content"))
166    split_stdout = stdout.splitlines()[-_MAX_REPORTED_ERROR_LINES::]
167    return "\n".join(split_stdout)
168
169def _FilterUnusedContent(content):
170    """Filter unused content from html.
171
172    Remove the html tags and style from content.
173
174    Args:
175        content: String, html content.
176
177    Returns:
178        String without html style or tags.
179    """
180    content = re.sub(_ERROR_MSG_TO_QUOTE_RE, "'", content)
181    content = re.sub(_ERROR_MSG_DEL_STYLE_RE, "", content, flags=re.DOTALL)
182    content = re.sub(_ERROR_MSG_DEL_TAGS_RE, "", content)
183    content = re.sub(r"\\n", " ", content)
184    return content
185
186
187def ShellCmdWithRetry(cmd, timeout=None, show_output=False,
188                      retry=_SSH_CMD_MAX_RETRY):
189    """Runs a shell command on remote device.
190
191    If the network is unstable and causes SSH connect fail, it will retry. When
192    it retry in a short time, you may encounter unstable network. We will use
193    the mechanism of RETRY_BACKOFF_FACTOR. The retry time for each failure is
194    times * retries.
195
196    Args:
197        cmd: String of the full SSH command to run, including the SSH binary and its arguments.
198        timeout: Optional integer, number of seconds to give.
199        show_output: Boolean, True to show command output in screen.
200        retry: Integer, the retry times.
201
202    Raises:
203        errors.DeviceConnectionError: For any non-zero return code of remote_cmd.
204        errors.LaunchCVDFail: Happened on launch_cvd with specific pattern of error message.
205        subprocess.CalledProcessError: The process exited with an error on the instance.
206    """
207    utils.RetryExceptionType(
208        exception_types=(errors.DeviceConnectionError,
209                         errors.LaunchCVDFail,
210                         subprocess.CalledProcessError),
211        max_retries=retry,
212        functor=_SshLogOutput,
213        sleep_multiplier=_SSH_CMD_RETRY_SLEEP,
214        retry_backoff_factor=utils.DEFAULT_RETRY_BACKOFF_FACTOR,
215        cmd=cmd,
216        timeout=timeout,
217        show_output=show_output)
218
219
220class IP():
221    """ A class that control the IP address."""
222    def __init__(self, external=None, internal=None, ip=None):
223        """Init for IP.
224            Args:
225                external: String, external ip.
226                internal: String, internal ip.
227                ip: String, default ip to set for either external and internal
228                if neither is set.
229        """
230        self.external = external or ip
231        self.internal = internal or ip
232
233
234class Ssh():
235    """A class that control the remote instance via the IP address.
236
237    Attributes:
238        _ip: an IP object.
239        _user: String of user login into the instance.
240        _ssh_private_key_path: Path to the private key file.
241        _extra_args_ssh_tunnel: String, extra args for ssh or scp.
242    """
243    def __init__(self, ip, user, ssh_private_key_path,
244                 extra_args_ssh_tunnel=None, report_internal_ip=False):
245        self._ip = ip.internal if report_internal_ip else ip.external
246        self._user = user
247        self._ssh_private_key_path = ssh_private_key_path
248        self._extra_args_ssh_tunnel = extra_args_ssh_tunnel
249
250    def Run(self, target_command, timeout=None, show_output=False,
251            retry=_SSH_CMD_MAX_RETRY):
252        """Run a shell command over SSH on a remote instance.
253
254        Example:
255            ssh:
256                base_cmd_list is ["ssh", "-i", "~/private_key_path" ,"-l" , "user", "1.1.1.1"]
257                target_command is "remote command"
258            scp:
259                base_cmd_list is ["scp", "-i", "~/private_key_path"]
260                target_command is "{src_file} {dst_file}"
261
262        Args:
263            target_command: String, text of command to run on the remote instance.
264            timeout: Integer, the maximum time to wait for the command to respond.
265            show_output: Boolean, True to show command output in screen.
266            retry: Integer, the retry times.
267        """
268        ShellCmdWithRetry(self.GetBaseCmd(constants.SSH_BIN) + " " + target_command,
269                          timeout,
270                          show_output,
271                          retry)
272
273    def GetBaseCmd(self, execute_bin):
274        """Get a base command over SSH on a remote instance.
275
276        Example:
277            execute bin is ssh:
278                ssh -i ~/private_key_path $extra_args -l user 1.1.1.1
279            execute bin is scp:
280                scp -i ~/private_key_path $extra_args
281
282        Args:
283            execute_bin: String, execute type, e.g. ssh or scp.
284
285        Returns:
286            Strings of base connection command.
287
288        Raises:
289            errors.UnknownType: Don't support the execute bin.
290        """
291        base_cmd = [utils.FindExecutable(execute_bin)]
292        base_cmd.append(_SSH_CMD % {"rsa_key_file": self._ssh_private_key_path})
293        if self._extra_args_ssh_tunnel:
294            base_cmd.append(self._extra_args_ssh_tunnel)
295
296        if execute_bin == constants.SSH_BIN:
297            base_cmd.append(_SSH_IDENTITY %
298                            {"login_user":self._user, "ip_addr":self._ip})
299            return " ".join(base_cmd)
300        if execute_bin == constants.SCP_BIN:
301            return " ".join(base_cmd)
302
303        raise errors.UnknownType("Don't support the execute bin %s." % execute_bin)
304
305    def GetCmdOutput(self, cmd):
306        """Runs a single SSH command and get its output.
307
308        Args:
309            cmd: String, text of command to run on the remote instance.
310
311        Returns:
312            String of the command output.
313        """
314        ssh_cmd = "exec " + self.GetBaseCmd(constants.SSH_BIN) + " " + cmd
315        logger.info("Running command \"%s\"", ssh_cmd)
316        process = subprocess.Popen(ssh_cmd, shell=True, stdin=None,
317                                   stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
318                                   universal_newlines=True)
319        stdout, _ = process.communicate()
320        return stdout
321
322    def CheckSshConnection(self, timeout):
323        """Run remote 'uptime' ssh command to check ssh connection.
324
325        Args:
326            timeout: Integer, the maximum time to wait for the command to respond.
327
328        Raises:
329            errors.DeviceConnectionError: Ssh isn't ready in the remote instance.
330        """
331        remote_cmd = [self.GetBaseCmd(constants.SSH_BIN)]
332        remote_cmd.append("uptime")
333
334        if _SshCallWait(" ".join(remote_cmd), timeout) == 0:
335            return
336        raise errors.DeviceConnectionError(
337            "Ssh isn't ready in the remote instance.")
338
339    @utils.TimeExecute(function_description="Waiting for SSH server")
340    def WaitForSsh(self, timeout=None, max_retry=_SSH_CMD_MAX_RETRY):
341        """Wait until the remote instance is ready to accept commands over SSH.
342
343        Args:
344            timeout: Integer, the maximum time in seconds to wait for the
345                     command to respond.
346            max_retry: Integer, the maximum number of retry.
347
348        Raises:
349            errors.DeviceConnectionError: Ssh isn't ready in the remote instance.
350        """
351        ssh_timeout = timeout or constants.DEFAULT_SSH_TIMEOUT
352        sleep_multiplier = ssh_timeout / sum(range(max_retry + 1))
353        logger.debug("Retry with interval time: %s secs", str(sleep_multiplier))
354        try:
355            utils.RetryExceptionType(
356                exception_types=errors.DeviceConnectionError,
357                max_retries=max_retry,
358                functor=self.CheckSshConnection,
359                sleep_multiplier=sleep_multiplier,
360                retry_backoff_factor=utils.DEFAULT_RETRY_BACKOFF_FACTOR,
361                timeout=_CONNECTION_TIMEOUT)
362        except errors.DeviceConnectionError as ssh_timeout:
363            ssh_cmd = "%s uptime" % self.GetBaseCmd(constants.SSH_BIN)
364            _SshLogOutput(ssh_cmd, timeout=_CONNECTION_TIMEOUT)
365            raise errors.DeviceConnectionError(
366                "Ssh connect timeout.\nYou can try the ssh connect command to "
367                "get detail information: '%s'" % ssh_cmd) from ssh_timeout
368
369    def ScpPushFile(self, src_file, dst_file):
370        """Scp push file to remote.
371
372        Args:
373            src_file: The source file path to be pulled.
374            dst_file: The destination file path the file is pulled to.
375        """
376        scp_command = [self.GetBaseCmd(constants.SCP_BIN)]
377        scp_command.append(src_file)
378        scp_command.append("%s@%s:%s" %(self._user, self._ip, dst_file))
379        ShellCmdWithRetry(" ".join(scp_command))
380
381    def ScpPushFiles(self, src_files, dst_dir):
382        """Push files to one specific folder of remote instance via scp command.
383
384        Args:
385            src_files: The source file path list to be pushed.
386            dst_dir: The destination directory the files to be pushed to.
387        """
388        scp_command = [self.GetBaseCmd(constants.SCP_BIN)]
389        scp_command.extend(src_files)
390        scp_command.append("%s@%s:%s" % (self._user, self._ip, dst_dir))
391        ShellCmdWithRetry(" ".join(scp_command))
392
393    def ScpPullFile(self, src_file, dst_file):
394        """Scp pull file from remote.
395
396        Args:
397            src_file: The source file path to be pulled.
398            dst_file: The destination file path the file is pulled to.
399        """
400        scp_command = [self.GetBaseCmd(constants.SCP_BIN)]
401        scp_command.append("%s@%s:%s" %(self._user, self._ip, src_file))
402        scp_command.append(dst_file)
403        ShellCmdWithRetry(" ".join(scp_command))
404