1#
2# Copyright 2007 Google Inc. Released under the GPL v2
3
4"""
5This module defines the SSHHost class.
6
7Implementation details:
8You should import the "hosts" package instead of importing each type of host.
9
10        SSHHost: a remote machine with a ssh access
11"""
12
13import re, logging
14from autotest_lib.client.common_lib import error, pxssh
15from autotest_lib.client.common_lib.cros.graphite import autotest_stats
16from autotest_lib.server import utils
17from autotest_lib.server.hosts import abstract_ssh
18
19
20class SSHHost(abstract_ssh.AbstractSSHHost):
21    """
22    This class represents a remote machine controlled through an ssh
23    session on which you can run programs.
24
25    It is not the machine autoserv is running on. The machine must be
26    configured for password-less login, for example through public key
27    authentication.
28
29    It includes support for controlling the machine through a serial
30    console on which you can run programs. If such a serial console is
31    set up on the machine then capabilities such as hard reset and
32    boot strap monitoring are available. If the machine does not have a
33    serial console available then ordinary SSH-based commands will
34    still be available, but attempts to use extensions such as
35    console logging or hard reset will fail silently.
36
37    Implementation details:
38    This is a leaf class in an abstract class hierarchy, it must
39    implement the unimplemented methods in parent classes.
40    """
41
42    def _initialize(self, hostname, *args, **dargs):
43        """
44        Construct a SSHHost object
45
46        Args:
47                hostname: network hostname or address of remote machine
48        """
49        super(SSHHost, self)._initialize(hostname=hostname, *args, **dargs)
50        self.setup_ssh()
51
52
53    def ssh_command(self, connect_timeout=30, options='', alive_interval=300):
54        """
55        Construct an ssh command with proper args for this host.
56
57        @param connect_timeout: connection timeout (in seconds)
58        @param options: SSH options
59        @param alive_interval: SSH Alive interval.
60
61        """
62        options = "%s %s" % (options, self.master_ssh_option)
63        base_cmd = self.make_ssh_command(user=self.user, port=self.port,
64                                         opts=options,
65                                         hosts_file=self.known_hosts_file,
66                                         connect_timeout=connect_timeout,
67                                         alive_interval=alive_interval)
68        return "%s %s" % (base_cmd, self.hostname)
69
70
71    def _run(self, command, timeout, ignore_status,
72             stdout, stderr, connect_timeout, env, options, stdin, args,
73             ignore_timeout):
74        """Helper function for run()."""
75        ssh_cmd = self.ssh_command(connect_timeout, options)
76        if not env.strip():
77            env = ""
78        else:
79            env = "export %s;" % env
80        for arg in args:
81            command += ' "%s"' % utils.sh_escape(arg)
82        full_cmd = '%s "%s %s"' % (ssh_cmd, env, utils.sh_escape(command))
83
84        # TODO(jrbarnette):  crbug.com/484726 - When we're in an SSP
85        # container, sometimes shortly after reboot we will see DNS
86        # resolution errors on ssh commands; the problem never
87        # occurs more than once in a row.  This especially affects
88        # the autoupdate_Rollback test, but other cases have been
89        # affected, too.
90        #
91        # We work around it by detecting the first DNS resolution error
92        # and retrying exactly one time.
93        dns_retry_count = 2
94        while True:
95            result = utils.run(full_cmd, timeout, True, stdout, stderr,
96                               verbose=False, stdin=stdin,
97                               stderr_is_expected=ignore_status,
98                               ignore_timeout=ignore_timeout)
99            dns_retry_count -= 1
100            if (result and result.exit_status == 255 and
101                    re.search(r'^ssh: .*: Name or service not known',
102                              result.stderr)):
103                if dns_retry_count:
104                    logging.debug('Retrying because of DNS failure')
105                    continue
106                logging.debug('Retry failed.')
107                autotest_stats.Counter('dns_retry_hack.fail').increment()
108            elif not dns_retry_count:
109                logging.debug('Retry succeeded.')
110                autotest_stats.Counter('dns_retry_hack.pass').increment()
111            break
112
113        if ignore_timeout and not result:
114            return None
115
116        # The error messages will show up in band (indistinguishable
117        # from stuff sent through the SSH connection), so we have the
118        # remote computer echo the message "Connected." before running
119        # any command.  Since the following 2 errors have to do with
120        # connecting, it's safe to do these checks.
121        if result.exit_status == 255:
122            if re.search(r'^ssh: connect to host .* port .*: '
123                         r'Connection timed out\r$', result.stderr):
124                raise error.AutoservSSHTimeout("ssh timed out", result)
125            if "Permission denied." in result.stderr:
126                msg = "ssh permission denied"
127                raise error.AutoservSshPermissionDeniedError(msg, result)
128
129        if not ignore_status and result.exit_status > 0:
130            raise error.AutoservRunError("command execution error", result)
131
132        return result
133
134
135    def run(self, command, timeout=3600, ignore_status=False,
136            stdout_tee=utils.TEE_TO_LOGS, stderr_tee=utils.TEE_TO_LOGS,
137            connect_timeout=30, options='', stdin=None, verbose=True, args=(),
138            ignore_timeout=False):
139        """
140        Run a command on the remote host.
141        @see common_lib.hosts.host.run()
142
143        @param connect_timeout: connection timeout (in seconds)
144        @param options: string with additional ssh command options
145        @param verbose: log the commands
146        @param ignore_timeout: bool True if SSH command timeouts should be
147                ignored.  Will return None on command timeout.
148
149        @raises AutoservRunError: if the command failed
150        @raises AutoservSSHTimeout: ssh connection has timed out
151        """
152        if verbose:
153            logging.debug("Running (ssh) '%s'", command)
154
155        # Start a master SSH connection if necessary.
156        self.start_master_ssh()
157
158        env = " ".join("=".join(pair) for pair in self.env.iteritems())
159        try:
160            return self._run(command, timeout, ignore_status,
161                             stdout_tee, stderr_tee, connect_timeout, env,
162                             options, stdin, args, ignore_timeout)
163        except error.CmdError, cmderr:
164            # We get a CmdError here only if there is timeout of that command.
165            # Catch that and stuff it into AutoservRunError and raise it.
166            timeout_message = str('Timeout encountered: %s' % cmderr.args[0])
167            raise error.AutoservRunError(timeout_message, cmderr.args[1])
168
169
170    def run_background(self, command, verbose=True):
171        """Start a command on the host in the background.
172
173        The command is started on the host in the background, and
174        this method call returns immediately without waiting for the
175        command's completion.  The PID of the process on the host is
176        returned as a string.
177
178        The command may redirect its stdin, stdout, or stderr as
179        necessary.  Without redirection, all input and output will
180        use /dev/null.
181
182        @param command The command to run in the background
183        @param verbose As for `self.run()`
184
185        @return Returns the PID of the remote background process
186                as a string.
187        """
188        # Redirection here isn't merely hygienic; it's a functional
189        # requirement.  sshd won't terminate until stdin, stdout,
190        # and stderr are all closed.
191        #
192        # The subshell is needed to do the right thing in case the
193        # passed in command has its own I/O redirections.
194        cmd_fmt = '( %s ) </dev/null >/dev/null 2>&1 & echo -n $!'
195        return self.run(cmd_fmt % command, verbose=verbose).stdout
196
197
198    def run_short(self, command, **kwargs):
199        """
200        Calls the run() command with a short default timeout.
201
202        Takes the same arguments as does run(),
203        with the exception of the timeout argument which
204        here is fixed at 60 seconds.
205        It returns the result of run.
206
207        @param command: the command line string
208
209        """
210        return self.run(command, timeout=60, **kwargs)
211
212
213    def run_grep(self, command, timeout=30, ignore_status=False,
214                 stdout_ok_regexp=None, stdout_err_regexp=None,
215                 stderr_ok_regexp=None, stderr_err_regexp=None,
216                 connect_timeout=30):
217        """
218        Run a command on the remote host and look for regexp
219        in stdout or stderr to determine if the command was
220        successul or not.
221
222
223        @param command: the command line string
224        @param timeout: time limit in seconds before attempting to
225                        kill the running process. The run() function
226                        will take a few seconds longer than 'timeout'
227                        to complete if it has to kill the process.
228        @param ignore_status: do not raise an exception, no matter
229                              what the exit code of the command is.
230        @param stdout_ok_regexp: regexp that should be in stdout
231                                 if the command was successul.
232        @param stdout_err_regexp: regexp that should be in stdout
233                                  if the command failed.
234        @param stderr_ok_regexp: regexp that should be in stderr
235                                 if the command was successul.
236        @param stderr_err_regexp: regexp that should be in stderr
237                                 if the command failed.
238        @param connect_timeout: connection timeout (in seconds)
239
240        Returns:
241                if the command was successul, raises an exception
242                otherwise.
243
244        Raises:
245                AutoservRunError:
246                - the exit code of the command execution was not 0.
247                - If stderr_err_regexp is found in stderr,
248                - If stdout_err_regexp is found in stdout,
249                - If stderr_ok_regexp is not found in stderr.
250                - If stdout_ok_regexp is not found in stdout,
251        """
252
253        # We ignore the status, because we will handle it at the end.
254        result = self.run(command, timeout, ignore_status=True,
255                          connect_timeout=connect_timeout)
256
257        # Look for the patterns, in order
258        for (regexp, stream) in ((stderr_err_regexp, result.stderr),
259                                 (stdout_err_regexp, result.stdout)):
260            if regexp and stream:
261                err_re = re.compile (regexp)
262                if err_re.search(stream):
263                    raise error.AutoservRunError(
264                        '%s failed, found error pattern: "%s"' % (command,
265                                                                regexp), result)
266
267        for (regexp, stream) in ((stderr_ok_regexp, result.stderr),
268                                 (stdout_ok_regexp, result.stdout)):
269            if regexp and stream:
270                ok_re = re.compile (regexp)
271                if ok_re.search(stream):
272                    if ok_re.search(stream):
273                        return
274
275        if not ignore_status and result.exit_status > 0:
276            raise error.AutoservRunError("command execution error", result)
277
278
279    def setup_ssh_key(self):
280        """Setup SSH Key"""
281        logging.debug('Performing SSH key setup on %s:%d as %s.',
282                      self.hostname, self.port, self.user)
283
284        try:
285            host = pxssh.pxssh()
286            host.login(self.hostname, self.user, self.password,
287                        port=self.port)
288            public_key = utils.get_public_key()
289
290            host.sendline('mkdir -p ~/.ssh')
291            host.prompt()
292            host.sendline('chmod 700 ~/.ssh')
293            host.prompt()
294            host.sendline("echo '%s' >> ~/.ssh/authorized_keys; " %
295                            public_key)
296            host.prompt()
297            host.sendline('chmod 600 ~/.ssh/authorized_keys')
298            host.prompt()
299            host.logout()
300
301            logging.debug('SSH key setup complete.')
302
303        except:
304            logging.debug('SSH key setup has failed.')
305            try:
306                host.logout()
307            except:
308                pass
309
310
311    def setup_ssh(self):
312        """Setup SSH"""
313        if self.password:
314            try:
315                self.ssh_ping()
316            except error.AutoservSshPingHostError:
317                self.setup_ssh_key()
318