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 self._os_if.log('Executing: %s' % cmd) 34 process = subprocess.Popen( 35 cmd, 36 shell=True, 37 stdout=subprocess.PIPE, 38 stderr=subprocess.PIPE) 39 if block: 40 process.wait() 41 return process 42 43 def run_command(self, cmd, block=True): 44 """Run a shell command. 45 46 In case of the command returning an error print its stdout and stderr 47 outputs on the console and dump them into the log. Otherwise suppress 48 all output. 49 50 @param block: if True (default), wait for command to finish 51 @raise error.CmdError: if block is True and command fails (rc!=0) 52 """ 53 start_time = time.time() 54 process = self._run_command(cmd, block) 55 if block and process.returncode: 56 # Grab output only if an error occurred 57 returncode = process.returncode 58 stdout = process.stdout.read() 59 stderr = process.stderr.read() 60 duration = time.time() - start_time 61 result = utils.CmdResult(cmd, stdout, stderr, returncode, duration) 62 self._os_if.log('Command failed.\n%s' % result) 63 raise error.CmdError(cmd, result) 64 65 def run_command_get_result(self, cmd, ignore_status=False): 66 """Run a shell command, and get the result (output and returncode). 67 68 @param ignore_status: if True, do not raise CmdError, even if rc != 0. 69 @raise error.CmdError: if command fails (rc!=0) and not ignore_result 70 @return the result of the command 71 @rtype: utils.CmdResult 72 """ 73 start_time = time.time() 74 75 process = self._run_command(cmd, block=True) 76 77 returncode = process.returncode 78 stdout = process.stdout.read() 79 stderr = process.stderr.read() 80 duration = time.time() - start_time 81 result = utils.CmdResult(cmd, stdout, stderr, returncode, duration) 82 83 if returncode and not ignore_status: 84 self._os_if.log('Command failed:\n%s' % result) 85 raise error.CmdError(cmd, result) 86 87 self._os_if.log('Command result:\n%s' % result) 88 return result 89 90 def run_command_check_output(self, cmd, success_token): 91 """Run a command and check whether standard output contains some string. 92 93 The sucess token is assumed to not contain newlines. 94 95 @param cmd: A string of the command to make a blocking call with. 96 @param success_token: A string to search the standard output of the 97 command for. 98 99 @returns a Boolean indicating whthere the success_token was in the 100 stdout of the cmd. 101 102 @raises UnsupportedSuccessToken if a newline is found in the 103 success_token. 104 """ 105 # The run_command_get_outuput method strips newlines from stdout. 106 if '\n' in success_token: 107 raise UnsupportedSuccessToken() 108 cmd_stdout = ''.join(self.run_command_get_output(cmd)) 109 self._os_if.log('Checking for %s in %s' % (success_token, cmd_stdout)) 110 return success_token in cmd_stdout 111 112 def run_command_get_status(self, cmd): 113 """Run a shell command and return its return code. 114 115 The return code of the command is returned, in case of any error. 116 """ 117 process = self._run_command(cmd) 118 return process.returncode 119 120 def run_command_get_output(self, cmd, include_stderr=False): 121 """Run shell command and return stdout (and possibly stderr) to the caller. 122 123 The output is returned as a list of strings stripped of the newline 124 characters. 125 """ 126 process = self._run_command(cmd) 127 text = [x.rstrip() for x in process.stdout.readlines()] 128 if include_stderr: 129 text.extend([x.rstrip() for x in process.stderr.readlines()]) 130 return text 131 132 def read_file(self, path): 133 """Read the content of the file.""" 134 with open(path) as f: 135 return f.read() 136 137 def write_file(self, path, data): 138 """Write the data to the file.""" 139 with open(path, 'w') as f: 140 f.write(data) 141 142 def append_file(self, path, data): 143 """Append the data to the file.""" 144 with open(path, 'a') as f: 145 f.write(data) 146