1# Copyright 2015 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4"""A module to abstract the shell execution environment on DUT."""
5
6import subprocess
7
8import time
9
10from autotest_lib.client.common_lib import error
11from autotest_lib.client.common_lib import utils
12
13
14class UnsupportedSuccessToken(Exception):
15    """Unsupported character found."""
16    pass
17
18
19class LocalShell(object):
20    """An object to wrap the local shell environment."""
21
22    def __init__(self, os_if):
23        """Initialize the LocalShell object."""
24        self._os_if = os_if
25
26    def _run_command(self, cmd, block=True):
27        """Helper function of run_command() methods.
28
29        Return the subprocess.Popen() instance to provide access to console
30        output in case command succeeded.  If block=False, will not wait for
31        process to return before returning.
32        """
33        stdout = None
34        stderr = None
35        if cmd and cmd.rstrip()[-1] == '&' and block:
36            errormsg = ('Remove & from command \'%s\', '
37                        'use block=True instead, '
38                        'refer to b/172325331 for more details' % cmd)
39            raise UnsupportedSuccessToken(errormsg)
40        self._os_if.log('Executing: %s' % cmd)
41        process = subprocess.Popen(
42                cmd,
43                shell=True,
44                stdout=subprocess.PIPE,
45                stderr=subprocess.PIPE)
46        if block:
47            stdout, stderr = process.communicate()
48        return process, stdout, stderr
49
50    def run_command(self, cmd, block=True):
51        """Run a shell command.
52
53        In case of the command returning an error print its stdout and stderr
54        outputs on the console and dump them into the log. Otherwise suppress
55        all output.
56
57        @param block: if True (default), wait for command to finish
58        @raise error.CmdError: if block is True and command fails (rc!=0)
59        """
60        start_time = time.time()
61        process, stdout, stderr = self._run_command(cmd, block)
62        if block and process.returncode:
63            # Grab output only if an error occurred
64            returncode = process.returncode
65            duration = time.time() - start_time
66            result = utils.CmdResult(cmd, stdout, stderr, returncode, duration)
67            self._os_if.log('Command failed.\n%s' % result)
68            raise error.CmdError(cmd, result)
69
70    def run_command_get_result(self, cmd, ignore_status=False):
71        """Run a shell command, and get the result (output and returncode).
72
73        @param ignore_status: if True, do not raise CmdError, even if rc != 0.
74        @raise error.CmdError: if command fails (rc!=0) and not ignore_result
75        @return the result of the command
76        @rtype: utils.CmdResult
77        """
78        start_time = time.time()
79
80        process, stdout, stderr = self._run_command(cmd, block=True)
81
82        returncode = process.returncode
83        duration = time.time() - start_time
84        result = utils.CmdResult(cmd, stdout, stderr, returncode, duration)
85
86        if returncode and not ignore_status:
87            self._os_if.log('Command failed:\n%s' % result)
88            raise error.CmdError(cmd, result)
89
90        self._os_if.log('Command result:\n%s' % result)
91        return result
92
93    def run_command_check_output(self, cmd, success_token):
94        """Run a command and check whether standard output contains some string.
95
96        The sucess token is assumed to not contain newlines.
97
98        @param cmd: A string of the command to make a blocking call with.
99        @param success_token: A string to search the standard output of the
100                command for.
101
102        @returns a Boolean indicating whthere the success_token was in the
103                stdout of the cmd.
104
105        @raises UnsupportedSuccessToken if a newline is found in the
106                success_token.
107        """
108        # The run_command_get_outuput method strips newlines from stdout.
109        if '\n' in success_token:
110            raise UnsupportedSuccessToken()
111        cmd_stdout = ''.join(self.run_command_get_output(cmd))
112        self._os_if.log('Checking for %s in %s' % (success_token, cmd_stdout))
113        return success_token in cmd_stdout
114
115    def run_command_get_status(self, cmd):
116        """Run a shell command and return its return code.
117
118        The return code of the command is returned, in case of any error.
119        """
120        process, stdout, stderr = self._run_command(cmd)
121        return process.returncode
122
123    def run_command_get_output(self, cmd, include_stderr=False):
124        """Run shell command and return stdout (and possibly stderr) to the caller.
125
126        The output is returned as a list of strings stripped of the newline
127        characters.
128        """
129        process, stdout, stderr = self._run_command(cmd)
130        text = [x.rstrip() for x in stdout.splitlines()]
131        if include_stderr:
132            text.extend([x.rstrip() for x in stderr.splitlines()])
133        return text
134
135    def read_file(self, path):
136        """Read the content of the file."""
137        with open(path) as f:
138            return f.read()
139
140    def write_file(self, path, data):
141        """Write the data to the file."""
142        with open(path, 'w') as f:
143            f.write(data)
144
145    def append_file(self, path, data):
146        """Append the data to the file."""
147        with open(path, 'a') as f:
148            f.write(data)
149