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 StringIO 13 import subprocess 14 import time 15 16 # fcntl is not available on Windows. 17 try: 18 import fcntl 19 except ImportError: 20 fcntl = None 21 22 23 def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): 24 return subprocess.Popen( 25 args=args, cwd=cwd, stdout=stdout, stderr=stderr, 26 shell=shell, close_fds=True, env=env, 27 preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL)) 28 29 30 def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): 31 pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, 32 env=env) 33 pipe.communicate() 34 return pipe.wait() 35 36 37 def RunCmd(args, cwd=None): 38 """Opens a subprocess to execute a program and returns its return value. 39 40 Args: 41 args: A string or a sequence of program arguments. The program to execute is 42 the string or the first item in the args sequence. 43 cwd: If not None, the subprocess's current directory will be changed to 44 |cwd| before it's executed. 45 46 Returns: 47 Return code from the command execution. 48 """ 49 logging.info(str(args) + ' ' + (cwd or '')) 50 return Call(args, cwd=cwd) 51 52 53 def GetCmdOutput(args, cwd=None, shell=False): 54 """Open a subprocess to execute a program and returns its output. 55 56 Args: 57 args: A string or a sequence of program arguments. The program to execute is 58 the string or the first item in the args sequence. 59 cwd: If not None, the subprocess's current directory will be changed to 60 |cwd| before it's executed. 61 shell: Whether to execute args as a shell command. 62 63 Returns: 64 Captures and returns the command's stdout. 65 Prints the command's stderr to logger (which defaults to stdout). 66 """ 67 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) 68 return output 69 70 71 def GetCmdStatusAndOutput(args, cwd=None, shell=False): 72 """Executes a subprocess and returns its exit code and output. 73 74 Args: 75 args: A string or a sequence of program arguments. The program to execute is 76 the string or the first item in the args sequence. 77 cwd: If not None, the subprocess's current directory will be changed to 78 |cwd| before it's executed. 79 shell: Whether to execute args as a shell command. 80 81 Returns: 82 The 2-tuple (exit code, output). 83 """ 84 if isinstance(args, basestring): 85 args_repr = args 86 if not shell: 87 raise Exception('string args must be run with shell=True') 88 elif shell: 89 raise Exception('array args must be run with shell=False') 90 else: 91 args_repr = ' '.join(map(pipes.quote, args)) 92 93 s = '[host]' 94 if cwd: 95 s += ':' + cwd 96 s += '> ' + args_repr 97 logging.info(s) 98 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 99 shell=shell, cwd=cwd) 100 stdout, stderr = pipe.communicate() 101 102 if stderr: 103 logging.critical(stderr) 104 if len(stdout) > 4096: 105 logging.debug('Truncated output:') 106 logging.debug(stdout[:4096]) 107 return (pipe.returncode, stdout) 108 109 110 class TimeoutError(Exception): 111 """Module-specific timeout exception.""" 112 pass 113 114 115 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, 116 logfile=None): 117 """Executes a subprocess with a timeout. 118 119 Args: 120 args: List of arguments to the program, the program to execute is the first 121 element. 122 timeout: the timeout in seconds or None to wait forever. 123 cwd: If not None, the subprocess's current directory will be changed to 124 |cwd| before it's executed. 125 shell: Whether to execute args as a shell command. 126 logfile: Optional file-like object that will receive output from the 127 command as it is running. 128 129 Returns: 130 The 2-tuple (exit code, output). 131 """ 132 assert fcntl, 'fcntl module is required' 133 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, 134 stderr=subprocess.STDOUT) 135 try: 136 end_time = (time.time() + timeout) if timeout else None 137 poll_interval = 1 138 buffer_size = 4096 139 child_fd = process.stdout.fileno() 140 output = StringIO.StringIO() 141 142 # Enable non-blocking reads from the child's stdout. 143 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) 144 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) 145 146 while True: 147 if end_time and time.time() > end_time: 148 raise TimeoutError 149 read_fds, _, _ = select.select([child_fd], [], [], poll_interval) 150 if child_fd in read_fds: 151 data = os.read(child_fd, buffer_size) 152 if not data: 153 break 154 if logfile: 155 logfile.write(data) 156 output.write(data) 157 if process.poll() is not None: 158 break 159 finally: 160 try: 161 # Make sure the process doesn't stick around if we fail with an 162 # exception. 163 process.kill() 164 except OSError: 165 pass 166 process.wait() 167 return process.returncode, output.getvalue() 168