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