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