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 19logger = logging.getLogger(__name__) 20 21_SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./') 22 23 24def 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 42def 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 70def 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 91def 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 105def 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 112def 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 128def 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 148def _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(str(c)) for c in args) 156 if cwd is None: 157 cwd = '' 158 else: 159 cwd = ':' + cwd 160 logger.debug('[host]%s> %s', cwd, args) 161 return args 162 163 164def 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 190def StartCmd(args, cwd=None, shell=False, env=None): 191 """Starts a subprocess and returns a handle to the process. 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 A process handle from subprocess.Popen. 205 """ 206 _ValidateAndLogCommand(args, cwd, shell) 207 return Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 208 shell=shell, cwd=cwd, env=env) 209 210 211def GetCmdStatusOutputAndError(args, cwd=None, shell=False, env=None): 212 """Executes a subprocess and returns its exit code, output, and errors. 213 214 Args: 215 args: A string or a sequence of program arguments. The program to execute is 216 the string or the first item in the args sequence. 217 cwd: If not None, the subprocess's current directory will be changed to 218 |cwd| before it's executed. 219 shell: Whether to execute args as a shell command. Must be True if args 220 is a string and False if args is a sequence. 221 env: If not None, a mapping that defines environment variables for the 222 subprocess. 223 224 Returns: 225 The 3-tuple (exit code, stdout, stderr). 226 """ 227 _ValidateAndLogCommand(args, cwd, shell) 228 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 229 shell=shell, cwd=cwd, env=env) 230 stdout, stderr = pipe.communicate() 231 return (pipe.returncode, stdout, stderr) 232 233 234class TimeoutError(Exception): 235 """Module-specific timeout exception.""" 236 237 def __init__(self, output=None): 238 super(TimeoutError, self).__init__() 239 self._output = output 240 241 @property 242 def output(self): 243 return self._output 244 245 246def _IterProcessStdoutFcntl( 247 process, iter_timeout=None, timeout=None, buffer_size=4096, 248 poll_interval=1): 249 """An fcntl-based implementation of _IterProcessStdout.""" 250 # pylint: disable=too-many-nested-blocks 251 import fcntl 252 try: 253 # Enable non-blocking reads from the child's stdout. 254 child_fd = process.stdout.fileno() 255 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) 256 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) 257 258 end_time = (time.time() + timeout) if timeout else None 259 iter_end_time = (time.time() + iter_timeout) if iter_timeout else None 260 261 while True: 262 if end_time and time.time() > end_time: 263 raise TimeoutError() 264 if iter_end_time and time.time() > iter_end_time: 265 yield None 266 iter_end_time = time.time() + iter_timeout 267 268 if iter_end_time: 269 iter_aware_poll_interval = min( 270 poll_interval, 271 max(0, iter_end_time - time.time())) 272 else: 273 iter_aware_poll_interval = poll_interval 274 275 read_fds, _, _ = select.select( 276 [child_fd], [], [], iter_aware_poll_interval) 277 if child_fd in read_fds: 278 data = os.read(child_fd, buffer_size) 279 if not data: 280 break 281 yield data 282 283 if process.poll() is not None: 284 # If process is closed, keep checking for output data (because of timing 285 # issues). 286 while True: 287 read_fds, _, _ = select.select( 288 [child_fd], [], [], iter_aware_poll_interval) 289 if child_fd in read_fds: 290 data = os.read(child_fd, buffer_size) 291 if data: 292 yield data 293 continue 294 break 295 break 296 finally: 297 try: 298 if process.returncode is None: 299 # Make sure the process doesn't stick around if we fail with an 300 # exception. 301 process.kill() 302 except OSError: 303 pass 304 process.wait() 305 306 307def _IterProcessStdoutQueue( 308 process, iter_timeout=None, timeout=None, buffer_size=4096, 309 poll_interval=1): 310 """A Queue.Queue-based implementation of _IterProcessStdout. 311 312 TODO(jbudorick): Evaluate whether this is a suitable replacement for 313 _IterProcessStdoutFcntl on all platforms. 314 """ 315 # pylint: disable=unused-argument 316 import Queue 317 import threading 318 319 stdout_queue = Queue.Queue() 320 321 def read_process_stdout(): 322 # TODO(jbudorick): Pick an appropriate read size here. 323 while True: 324 try: 325 output_chunk = os.read(process.stdout.fileno(), buffer_size) 326 except IOError: 327 break 328 stdout_queue.put(output_chunk, True) 329 if not output_chunk and process.poll() is not None: 330 break 331 332 reader_thread = threading.Thread(target=read_process_stdout) 333 reader_thread.start() 334 335 end_time = (time.time() + timeout) if timeout else None 336 337 try: 338 while True: 339 if end_time and time.time() > end_time: 340 raise TimeoutError() 341 try: 342 s = stdout_queue.get(True, iter_timeout) 343 if not s: 344 break 345 yield s 346 except Queue.Empty: 347 yield None 348 finally: 349 try: 350 if process.returncode is None: 351 # Make sure the process doesn't stick around if we fail with an 352 # exception. 353 process.kill() 354 except OSError: 355 pass 356 process.wait() 357 reader_thread.join() 358 359 360_IterProcessStdout = ( 361 _IterProcessStdoutQueue 362 if sys.platform == 'win32' 363 else _IterProcessStdoutFcntl) 364"""Iterate over a process's stdout. 365 366This is intentionally not public. 367 368Args: 369 process: The process in question. 370 iter_timeout: An optional length of time, in seconds, to wait in 371 between each iteration. If no output is received in the given 372 time, this generator will yield None. 373 timeout: An optional length of time, in seconds, during which 374 the process must finish. If it fails to do so, a TimeoutError 375 will be raised. 376 buffer_size: The maximum number of bytes to read (and thus yield) at once. 377 poll_interval: The length of time to wait in calls to `select.select`. 378 If iter_timeout is set, the remaining length of time in the iteration 379 may take precedence. 380Raises: 381 TimeoutError: if timeout is set and the process does not complete. 382Yields: 383 basestrings of data or None. 384""" 385 386 387def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, 388 logfile=None, env=None): 389 """Executes a subprocess with a timeout. 390 391 Args: 392 args: List of arguments to the program, the program to execute is the first 393 element. 394 timeout: the timeout in seconds or None to wait forever. 395 cwd: If not None, the subprocess's current directory will be changed to 396 |cwd| before it's executed. 397 shell: Whether to execute args as a shell command. Must be True if args 398 is a string and False if args is a sequence. 399 logfile: Optional file-like object that will receive output from the 400 command as it is running. 401 env: If not None, a mapping that defines environment variables for the 402 subprocess. 403 404 Returns: 405 The 2-tuple (exit code, output). 406 Raises: 407 TimeoutError on timeout. 408 """ 409 _ValidateAndLogCommand(args, cwd, shell) 410 output = StringIO.StringIO() 411 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, 412 stderr=subprocess.STDOUT, env=env) 413 try: 414 for data in _IterProcessStdout(process, timeout=timeout): 415 if logfile: 416 logfile.write(data) 417 output.write(data) 418 except TimeoutError: 419 raise TimeoutError(output.getvalue()) 420 421 str_output = output.getvalue() 422 logger.debug('STDOUT+STDERR: %s%s', str_output[:4096].rstrip(), 423 '<truncated>' if len(str_output) > 4096 else '') 424 return process.returncode, str_output 425 426 427def IterCmdOutputLines(args, iter_timeout=None, timeout=None, cwd=None, 428 shell=False, env=None, check_status=True): 429 """Executes a subprocess and continuously yields lines from its output. 430 431 Args: 432 args: List of arguments to the program, the program to execute is the first 433 element. 434 iter_timeout: Timeout for each iteration, in seconds. 435 timeout: Timeout for the entire command, in seconds. 436 cwd: If not None, the subprocess's current directory will be changed to 437 |cwd| before it's executed. 438 shell: Whether to execute args as a shell command. Must be True if args 439 is a string and False if args is a sequence. 440 env: If not None, a mapping that defines environment variables for the 441 subprocess. 442 check_status: A boolean indicating whether to check the exit status of the 443 process after all output has been read. 444 Yields: 445 The output of the subprocess, line by line. 446 447 Raises: 448 CalledProcessError if check_status is True and the process exited with a 449 non-zero exit status. 450 """ 451 cmd = _ValidateAndLogCommand(args, cwd, shell) 452 process = Popen(args, cwd=cwd, shell=shell, env=env, 453 stdout=subprocess.PIPE, stderr=subprocess.STDOUT) 454 return _IterCmdOutputLines( 455 process, cmd, iter_timeout=iter_timeout, timeout=timeout, 456 check_status=check_status) 457 458def _IterCmdOutputLines(process, cmd, iter_timeout=None, timeout=None, 459 check_status=True): 460 buffer_output = '' 461 462 iter_end = None 463 cur_iter_timeout = None 464 if iter_timeout: 465 iter_end = time.time() + iter_timeout 466 cur_iter_timeout = iter_timeout 467 468 for data in _IterProcessStdout(process, iter_timeout=cur_iter_timeout, 469 timeout=timeout): 470 if iter_timeout: 471 # Check whether the current iteration has timed out. 472 cur_iter_timeout = iter_end - time.time() 473 if data is None or cur_iter_timeout < 0: 474 yield None 475 iter_end = time.time() + iter_timeout 476 continue 477 else: 478 assert data is not None, ( 479 'Iteration received no data despite no iter_timeout being set. ' 480 'cmd: %s' % cmd) 481 482 # Construct lines to yield from raw data. 483 buffer_output += data 484 has_incomplete_line = buffer_output[-1] not in '\r\n' 485 lines = buffer_output.splitlines() 486 buffer_output = lines.pop() if has_incomplete_line else '' 487 for line in lines: 488 yield line 489 if iter_timeout: 490 iter_end = time.time() + iter_timeout 491 492 if buffer_output: 493 yield buffer_output 494 if check_status and process.returncode: 495 raise subprocess.CalledProcessError(process.returncode, cmd) 496