1# Copyright (c) 2012 The Chromium 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 5"""A wrapper for subprocess to make calling shell commands easier.""" 6 7import logging 8import os 9import pipes 10import select 11import signal 12import string 13import StringIO 14import subprocess 15import time 16 17# fcntl is not available on Windows. 18try: 19 import fcntl 20except ImportError: 21 fcntl = None 22 23_SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./') 24 25 26def SingleQuote(s): 27 """Return an shell-escaped version of the string using single quotes. 28 29 Reliably quote a string which may contain unsafe characters (e.g. space, 30 quote, or other special characters such as '$'). 31 32 The returned value can be used in a shell command line as one token that gets 33 to be interpreted literally. 34 35 Args: 36 s: The string to quote. 37 38 Return: 39 The string quoted using single quotes. 40 """ 41 return pipes.quote(s) 42 43 44def DoubleQuote(s): 45 """Return an shell-escaped version of the string using double quotes. 46 47 Reliably quote a string which may contain unsafe characters (e.g. space 48 or quote characters), while retaining some shell features such as variable 49 interpolation. 50 51 The returned value can be used in a shell command line as one token that gets 52 to be further interpreted by the shell. 53 54 The set of characters that retain their special meaning may depend on the 55 shell implementation. This set usually includes: '$', '`', '\', '!', '*', 56 and '@'. 57 58 Args: 59 s: The string to quote. 60 61 Return: 62 The string quoted using double quotes. 63 """ 64 if not s: 65 return '""' 66 elif all(c in _SafeShellChars for c in s): 67 return s 68 else: 69 return '"' + s.replace('"', '\\"') + '"' 70 71 72def ShrinkToSnippet(cmd_parts, var_name, var_value): 73 """Constructs a shell snippet for a command using a variable to shrink it. 74 75 Takes into account all quoting that needs to happen. 76 77 Args: 78 cmd_parts: A list of command arguments. 79 var_name: The variable that holds var_value. 80 var_value: The string to replace in cmd_parts with $var_name 81 82 Returns: 83 A shell snippet that does not include setting the variable. 84 """ 85 def shrink(value): 86 parts = (x and SingleQuote(x) for x in value.split(var_value)) 87 with_substitutions = ('"$%s"' % var_name).join(parts) 88 return with_substitutions or "''" 89 90 return ' '.join(shrink(part) for part in cmd_parts) 91 92 93def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): 94 return subprocess.Popen( 95 args=args, cwd=cwd, stdout=stdout, stderr=stderr, 96 shell=shell, close_fds=True, env=env, 97 preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL)) 98 99 100def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): 101 pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, 102 env=env) 103 pipe.communicate() 104 return pipe.wait() 105 106 107def RunCmd(args, cwd=None): 108 """Opens a subprocess to execute a program and returns its return value. 109 110 Args: 111 args: A string or a sequence of program arguments. The program to execute is 112 the string or the first item in the args sequence. 113 cwd: If not None, the subprocess's current directory will be changed to 114 |cwd| before it's executed. 115 116 Returns: 117 Return code from the command execution. 118 """ 119 logging.info(str(args) + ' ' + (cwd or '')) 120 return Call(args, cwd=cwd) 121 122 123def GetCmdOutput(args, cwd=None, shell=False): 124 """Open a subprocess to execute a program and returns its output. 125 126 Args: 127 args: A string or a sequence of program arguments. The program to execute is 128 the string or the first item in the args sequence. 129 cwd: If not None, the subprocess's current directory will be changed to 130 |cwd| before it's executed. 131 shell: Whether to execute args as a shell command. 132 133 Returns: 134 Captures and returns the command's stdout. 135 Prints the command's stderr to logger (which defaults to stdout). 136 """ 137 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) 138 return output 139 140 141def _ValidateAndLogCommand(args, cwd, shell): 142 if isinstance(args, basestring): 143 if not shell: 144 raise Exception('string args must be run with shell=True') 145 else: 146 if shell: 147 raise Exception('array args must be run with shell=False') 148 args = ' '.join(SingleQuote(c) for c in args) 149 if cwd is None: 150 cwd = '' 151 else: 152 cwd = ':' + cwd 153 logging.info('[host]%s> %s', cwd, args) 154 return args 155 156 157def GetCmdStatusAndOutput(args, cwd=None, shell=False): 158 """Executes a subprocess and returns its exit code and output. 159 160 Args: 161 args: A string or a sequence of program arguments. The program to execute is 162 the string or the first item in the args sequence. 163 cwd: If not None, the subprocess's current directory will be changed to 164 |cwd| before it's executed. 165 shell: Whether to execute args as a shell command. Must be True if args 166 is a string and False if args is a sequence. 167 168 Returns: 169 The 2-tuple (exit code, output). 170 """ 171 status, stdout, stderr = GetCmdStatusOutputAndError( 172 args, cwd=cwd, shell=shell) 173 174 if stderr: 175 logging.critical('STDERR: %s', stderr) 176 logging.debug('STDOUT: %s%s', stdout[:4096].rstrip(), 177 '<truncated>' if len(stdout) > 4096 else '') 178 return (status, stdout) 179 180 181def GetCmdStatusOutputAndError(args, cwd=None, shell=False): 182 """Executes a subprocess and returns its exit code, output, and errors. 183 184 Args: 185 args: A string or a sequence of program arguments. The program to execute is 186 the string or the first item in the args sequence. 187 cwd: If not None, the subprocess's current directory will be changed to 188 |cwd| before it's executed. 189 shell: Whether to execute args as a shell command. Must be True if args 190 is a string and False if args is a sequence. 191 192 Returns: 193 The 2-tuple (exit code, output). 194 """ 195 _ValidateAndLogCommand(args, cwd, shell) 196 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 197 shell=shell, cwd=cwd) 198 stdout, stderr = pipe.communicate() 199 return (pipe.returncode, stdout, stderr) 200 201 202class TimeoutError(Exception): 203 """Module-specific timeout exception.""" 204 205 def __init__(self, output=None): 206 super(TimeoutError, self).__init__() 207 self._output = output 208 209 @property 210 def output(self): 211 return self._output 212 213 214def _IterProcessStdout(process, timeout=None, buffer_size=4096, 215 poll_interval=1): 216 assert fcntl, 'fcntl module is required' 217 try: 218 # Enable non-blocking reads from the child's stdout. 219 child_fd = process.stdout.fileno() 220 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) 221 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) 222 223 end_time = (time.time() + timeout) if timeout else None 224 while True: 225 if end_time and time.time() > end_time: 226 raise TimeoutError() 227 read_fds, _, _ = select.select([child_fd], [], [], poll_interval) 228 if child_fd in read_fds: 229 data = os.read(child_fd, buffer_size) 230 if not data: 231 break 232 yield data 233 if process.poll() is not None: 234 break 235 finally: 236 try: 237 # Make sure the process doesn't stick around if we fail with an 238 # exception. 239 process.kill() 240 except OSError: 241 pass 242 process.wait() 243 244 245def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, 246 logfile=None): 247 """Executes a subprocess with a timeout. 248 249 Args: 250 args: List of arguments to the program, the program to execute is the first 251 element. 252 timeout: the timeout in seconds or None to wait forever. 253 cwd: If not None, the subprocess's current directory will be changed to 254 |cwd| before it's executed. 255 shell: Whether to execute args as a shell command. Must be True if args 256 is a string and False if args is a sequence. 257 logfile: Optional file-like object that will receive output from the 258 command as it is running. 259 260 Returns: 261 The 2-tuple (exit code, output). 262 """ 263 _ValidateAndLogCommand(args, cwd, shell) 264 output = StringIO.StringIO() 265 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, 266 stderr=subprocess.STDOUT) 267 try: 268 for data in _IterProcessStdout(process, timeout=timeout): 269 if logfile: 270 logfile.write(data) 271 output.write(data) 272 except TimeoutError: 273 raise TimeoutError(output.getvalue()) 274 275 str_output = output.getvalue() 276 logging.debug('STDOUT+STDERR: %s%s', str_output[:4096].rstrip(), 277 '<truncated>' if len(str_output) > 4096 else '') 278 return process.returncode, str_output 279 280 281def IterCmdOutputLines(args, timeout=None, cwd=None, shell=False, 282 check_status=True): 283 """Executes a subprocess and continuously yields lines from its output. 284 285 Args: 286 args: List of arguments to the program, the program to execute is the first 287 element. 288 cwd: If not None, the subprocess's current directory will be changed to 289 |cwd| before it's executed. 290 shell: Whether to execute args as a shell command. Must be True if args 291 is a string and False if args is a sequence. 292 check_status: A boolean indicating whether to check the exit status of the 293 process after all output has been read. 294 295 Yields: 296 The output of the subprocess, line by line. 297 298 Raises: 299 CalledProcessError if check_status is True and the process exited with a 300 non-zero exit status. 301 """ 302 cmd = _ValidateAndLogCommand(args, cwd, shell) 303 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, 304 stderr=subprocess.STDOUT) 305 buffer_output = '' 306 for data in _IterProcessStdout(process, timeout=timeout): 307 buffer_output += data 308 has_incomplete_line = buffer_output[-1] not in '\r\n' 309 lines = buffer_output.splitlines() 310 buffer_output = lines.pop() if has_incomplete_line else '' 311 for line in lines: 312 yield line 313 if buffer_output: 314 yield buffer_output 315 if check_status and process.returncode: 316 raise subprocess.CalledProcessError(process.returncode, cmd) 317