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