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