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 sys 16import time 17 18# fcntl is not available on Windows. 19try: 20 import fcntl 21except ImportError: 22 fcntl = None 23 24logger = logging.getLogger(__name__) 25 26_SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./') 27 28 29def SingleQuote(s): 30 """Return an shell-escaped version of the string using single quotes. 31 32 Reliably quote a string which may contain unsafe characters (e.g. space, 33 quote, or other special characters such as '$'). 34 35 The returned value can be used in a shell command line as one token that gets 36 to be interpreted literally. 37 38 Args: 39 s: The string to quote. 40 41 Return: 42 The string quoted using single quotes. 43 """ 44 return pipes.quote(s) 45 46 47def DoubleQuote(s): 48 """Return an shell-escaped version of the string using double quotes. 49 50 Reliably quote a string which may contain unsafe characters (e.g. space 51 or quote characters), while retaining some shell features such as variable 52 interpolation. 53 54 The returned value can be used in a shell command line as one token that gets 55 to be further interpreted by the shell. 56 57 The set of characters that retain their special meaning may depend on the 58 shell implementation. This set usually includes: '$', '`', '\', '!', '*', 59 and '@'. 60 61 Args: 62 s: The string to quote. 63 64 Return: 65 The string quoted using double quotes. 66 """ 67 if not s: 68 return '""' 69 elif all(c in _SafeShellChars for c in s): 70 return s 71 else: 72 return '"' + s.replace('"', '\\"') + '"' 73 74 75def ShrinkToSnippet(cmd_parts, var_name, var_value): 76 """Constructs a shell snippet for a command using a variable to shrink it. 77 78 Takes into account all quoting that needs to happen. 79 80 Args: 81 cmd_parts: A list of command arguments. 82 var_name: The variable that holds var_value. 83 var_value: The string to replace in cmd_parts with $var_name 84 85 Returns: 86 A shell snippet that does not include setting the variable. 87 """ 88 def shrink(value): 89 parts = (x and SingleQuote(x) for x in value.split(var_value)) 90 with_substitutions = ('"$%s"' % var_name).join(parts) 91 return with_substitutions or "''" 92 93 return ' '.join(shrink(part) for part in cmd_parts) 94 95 96def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): 97 # preexec_fn isn't supported on windows. 98 if sys.platform == 'win32': 99 preexec_fn = None 100 else: 101 preexec_fn = lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL) 102 103 return subprocess.Popen( 104 args=args, cwd=cwd, stdout=stdout, stderr=stderr, 105 shell=shell, close_fds=True, env=env, preexec_fn=preexec_fn) 106 107 108def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): 109 pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, 110 env=env) 111 pipe.communicate() 112 return pipe.wait() 113 114 115def RunCmd(args, cwd=None): 116 """Opens a subprocess to execute a program and returns its return value. 117 118 Args: 119 args: A string or a sequence of program arguments. The program to execute is 120 the string or the first item in the args sequence. 121 cwd: If not None, the subprocess's current directory will be changed to 122 |cwd| before it's executed. 123 124 Returns: 125 Return code from the command execution. 126 """ 127 logger.info(str(args) + ' ' + (cwd or '')) 128 return Call(args, cwd=cwd) 129 130 131def GetCmdOutput(args, cwd=None, shell=False): 132 """Open a subprocess to execute a program and returns its output. 133 134 Args: 135 args: A string or a sequence of program arguments. The program to execute is 136 the string or the first item in the args sequence. 137 cwd: If not None, the subprocess's current directory will be changed to 138 |cwd| before it's executed. 139 shell: Whether to execute args as a shell command. 140 141 Returns: 142 Captures and returns the command's stdout. 143 Prints the command's stderr to logger (which defaults to stdout). 144 """ 145 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) 146 return output 147 148 149def _ValidateAndLogCommand(args, cwd, shell): 150 if isinstance(args, basestring): 151 if not shell: 152 raise Exception('string args must be run with shell=True') 153 else: 154 if shell: 155 raise Exception('array args must be run with shell=False') 156 args = ' '.join(SingleQuote(c) for c in args) 157 if cwd is None: 158 cwd = '' 159 else: 160 cwd = ':' + cwd 161 logger.info('[host]%s> %s', cwd, args) 162 return args 163 164 165def GetCmdStatusAndOutput(args, cwd=None, shell=False): 166 """Executes a subprocess and returns its exit code and output. 167 168 Args: 169 args: A string or a sequence of program arguments. The program to execute is 170 the string or the first item in the args sequence. 171 cwd: If not None, the subprocess's current directory will be changed to 172 |cwd| before it's executed. 173 shell: Whether to execute args as a shell command. Must be True if args 174 is a string and False if args is a sequence. 175 176 Returns: 177 The 2-tuple (exit code, output). 178 """ 179 status, stdout, stderr = GetCmdStatusOutputAndError( 180 args, cwd=cwd, shell=shell) 181 182 if stderr: 183 logger.critical('STDERR: %s', stderr) 184 logger.debug('STDOUT: %s%s', stdout[:4096].rstrip(), 185 '<truncated>' if len(stdout) > 4096 else '') 186 return (status, stdout) 187 188 189def GetCmdStatusOutputAndError(args, cwd=None, shell=False): 190 """Executes a subprocess and returns its exit code, output, and errors. 191 192 Args: 193 args: A string or a sequence of program arguments. The program to execute is 194 the string or the first item in the args sequence. 195 cwd: If not None, the subprocess's current directory will be changed to 196 |cwd| before it's executed. 197 shell: Whether to execute args as a shell command. Must be True if args 198 is a string and False if args is a sequence. 199 200 Returns: 201 The 2-tuple (exit code, output). 202 """ 203 _ValidateAndLogCommand(args, cwd, shell) 204 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 205 shell=shell, cwd=cwd) 206 stdout, stderr = pipe.communicate() 207 return (pipe.returncode, stdout, stderr) 208 209 210class TimeoutError(Exception): 211 """Module-specific timeout exception.""" 212 213 def __init__(self, output=None): 214 super(TimeoutError, self).__init__() 215 self._output = output 216 217 @property 218 def output(self): 219 return self._output 220 221 222def _IterProcessStdout(process, iter_timeout=None, timeout=None, 223 buffer_size=4096, poll_interval=1): 224 """Iterate over a process's stdout. 225 226 This is intentionally not public. 227 228 Args: 229 process: The process in question. 230 iter_timeout: An optional length of time, in seconds, to wait in 231 between each iteration. If no output is received in the given 232 time, this generator will yield None. 233 timeout: An optional length of time, in seconds, during which 234 the process must finish. If it fails to do so, a TimeoutError 235 will be raised. 236 buffer_size: The maximum number of bytes to read (and thus yield) at once. 237 poll_interval: The length of time to wait in calls to `select.select`. 238 If iter_timeout is set, the remaining length of time in the iteration 239 may take precedence. 240 Raises: 241 TimeoutError: if timeout is set and the process does not complete. 242 Yields: 243 basestrings of data or None. 244 """ 245 246 assert fcntl, 'fcntl module is required' 247 try: 248 # Enable non-blocking reads from the child's stdout. 249 child_fd = process.stdout.fileno() 250 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) 251 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) 252 253 end_time = (time.time() + timeout) if timeout else None 254 iter_end_time = (time.time() + iter_timeout) if iter_timeout else None 255 256 while True: 257 if end_time and time.time() > end_time: 258 raise TimeoutError() 259 if iter_end_time and time.time() > iter_end_time: 260 yield None 261 iter_end_time = time.time() + iter_timeout 262 263 if iter_end_time: 264 iter_aware_poll_interval = min( 265 poll_interval, 266 max(0, iter_end_time - time.time())) 267 else: 268 iter_aware_poll_interval = poll_interval 269 270 read_fds, _, _ = select.select( 271 [child_fd], [], [], iter_aware_poll_interval) 272 if child_fd in read_fds: 273 data = os.read(child_fd, buffer_size) 274 if not data: 275 break 276 yield data 277 if process.poll() is not None: 278 break 279 finally: 280 try: 281 if process.returncode is None: 282 # Make sure the process doesn't stick around if we fail with an 283 # exception. 284 process.kill() 285 except OSError: 286 pass 287 process.wait() 288 289 290def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, 291 logfile=None): 292 """Executes a subprocess with a timeout. 293 294 Args: 295 args: List of arguments to the program, the program to execute is the first 296 element. 297 timeout: the timeout in seconds or None to wait forever. 298 cwd: If not None, the subprocess's current directory will be changed to 299 |cwd| before it's executed. 300 shell: Whether to execute args as a shell command. Must be True if args 301 is a string and False if args is a sequence. 302 logfile: Optional file-like object that will receive output from the 303 command as it is running. 304 305 Returns: 306 The 2-tuple (exit code, output). 307 Raises: 308 TimeoutError on timeout. 309 """ 310 _ValidateAndLogCommand(args, cwd, shell) 311 output = StringIO.StringIO() 312 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, 313 stderr=subprocess.STDOUT) 314 try: 315 for data in _IterProcessStdout(process, timeout=timeout): 316 if logfile: 317 logfile.write(data) 318 output.write(data) 319 except TimeoutError: 320 raise TimeoutError(output.getvalue()) 321 322 str_output = output.getvalue() 323 logger.debug('STDOUT+STDERR: %s%s', str_output[:4096].rstrip(), 324 '<truncated>' if len(str_output) > 4096 else '') 325 return process.returncode, str_output 326 327 328def IterCmdOutputLines(args, iter_timeout=None, timeout=None, cwd=None, 329 shell=False, check_status=True): 330 """Executes a subprocess and continuously yields lines from its output. 331 332 Args: 333 args: List of arguments to the program, the program to execute is the first 334 element. 335 iter_timeout: Timeout for each iteration, in seconds. 336 timeout: Timeout for the entire command, in seconds. 337 cwd: If not None, the subprocess's current directory will be changed to 338 |cwd| before it's executed. 339 shell: Whether to execute args as a shell command. Must be True if args 340 is a string and False if args is a sequence. 341 check_status: A boolean indicating whether to check the exit status of the 342 process after all output has been read. 343 Yields: 344 The output of the subprocess, line by line. 345 346 Raises: 347 CalledProcessError if check_status is True and the process exited with a 348 non-zero exit status. 349 """ 350 cmd = _ValidateAndLogCommand(args, cwd, shell) 351 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, 352 stderr=subprocess.STDOUT) 353 return _IterCmdOutputLines( 354 process, cmd, iter_timeout=iter_timeout, timeout=timeout, 355 check_status=check_status) 356 357def _IterCmdOutputLines(process, cmd, iter_timeout=None, timeout=None, 358 check_status=True): 359 buffer_output = '' 360 361 iter_end = None 362 cur_iter_timeout = None 363 if iter_timeout: 364 iter_end = time.time() + iter_timeout 365 cur_iter_timeout = iter_timeout 366 367 for data in _IterProcessStdout(process, iter_timeout=cur_iter_timeout, 368 timeout=timeout): 369 if iter_timeout: 370 # Check whether the current iteration has timed out. 371 cur_iter_timeout = iter_end - time.time() 372 if data is None or cur_iter_timeout < 0: 373 yield None 374 iter_end = time.time() + iter_timeout 375 continue 376 else: 377 assert data is not None, ( 378 'Iteration received no data despite no iter_timeout being set. ' 379 'cmd: %s' % cmd) 380 381 # Construct lines to yield from raw data. 382 buffer_output += data 383 has_incomplete_line = buffer_output[-1] not in '\r\n' 384 lines = buffer_output.splitlines() 385 buffer_output = lines.pop() if has_incomplete_line else '' 386 for line in lines: 387 yield line 388 if iter_timeout: 389 iter_end = time.time() + iter_timeout 390 391 if buffer_output: 392 yield buffer_output 393 if check_status and process.returncode: 394 raise subprocess.CalledProcessError(process.returncode, cmd) 395