OLD | NEW |
| (Empty) |
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 time | |
16 | |
17 # fcntl is not available on Windows. | |
18 try: | |
19 import fcntl | |
20 except ImportError: | |
21 fcntl = None | |
22 | |
23 _SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./') | |
24 | |
25 def SingleQuote(s): | |
26 """Return an shell-escaped version of the string using single quotes. | |
27 | |
28 Reliably quote a string which may contain unsafe characters (e.g. space, | |
29 quote, or other special characters such as '$'). | |
30 | |
31 The returned value can be used in a shell command line as one token that gets | |
32 to be interpreted literally. | |
33 | |
34 Args: | |
35 s: The string to quote. | |
36 | |
37 Return: | |
38 The string quoted using single quotes. | |
39 """ | |
40 return pipes.quote(s) | |
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 return subprocess.Popen( | |
93 args=args, cwd=cwd, stdout=stdout, stderr=stderr, | |
94 shell=shell, close_fds=True, env=env, | |
95 preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL)) | |
96 | |
97 | |
98 def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None): | |
99 pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, | |
100 env=env) | |
101 pipe.communicate() | |
102 return pipe.wait() | |
103 | |
104 | |
105 def RunCmd(args, cwd=None): | |
106 """Opens a subprocess to execute a program and returns its return value. | |
107 | |
108 Args: | |
109 args: A string or a sequence of program arguments. The program to execute is | |
110 the string or the first item in the args sequence. | |
111 cwd: If not None, the subprocess's current directory will be changed to | |
112 |cwd| before it's executed. | |
113 | |
114 Returns: | |
115 Return code from the command execution. | |
116 """ | |
117 logging.info(str(args) + ' ' + (cwd or '')) | |
118 return Call(args, cwd=cwd) | |
119 | |
120 | |
121 def GetCmdOutput(args, cwd=None, shell=False): | |
122 """Open a subprocess to execute a program and returns its output. | |
123 | |
124 Args: | |
125 args: A string or a sequence of program arguments. The program to execute is | |
126 the string or the first item in the args sequence. | |
127 cwd: If not None, the subprocess's current directory will be changed to | |
128 |cwd| before it's executed. | |
129 shell: Whether to execute args as a shell command. | |
130 | |
131 Returns: | |
132 Captures and returns the command's stdout. | |
133 Prints the command's stderr to logger (which defaults to stdout). | |
134 """ | |
135 (_, output) = GetCmdStatusAndOutput(args, cwd, shell) | |
136 return output | |
137 | |
138 | |
139 def _ValidateAndLogCommand(args, cwd, shell): | |
140 if isinstance(args, basestring): | |
141 if not shell: | |
142 raise Exception('string args must be run with shell=True') | |
143 else: | |
144 if shell: | |
145 raise Exception('array args must be run with shell=False') | |
146 args = ' '.join(SingleQuote(c) for c in args) | |
147 if cwd is None: | |
148 cwd = '' | |
149 else: | |
150 cwd = ':' + cwd | |
151 logging.info('[host]%s> %s', cwd, args) | |
152 return args | |
153 | |
154 | |
155 def GetCmdStatusAndOutput(args, cwd=None, shell=False): | |
156 """Executes a subprocess and returns its exit code and output. | |
157 | |
158 Args: | |
159 args: A string or a sequence of program arguments. The program to execute is | |
160 the string or the first item in the args sequence. | |
161 cwd: If not None, the subprocess's current directory will be changed to | |
162 |cwd| before it's executed. | |
163 shell: Whether to execute args as a shell command. Must be True if args | |
164 is a string and False if args is a sequence. | |
165 | |
166 Returns: | |
167 The 2-tuple (exit code, output). | |
168 """ | |
169 status, stdout, stderr = GetCmdStatusOutputAndError( | |
170 args, cwd=cwd, shell=shell) | |
171 | |
172 if stderr: | |
173 logging.critical(stderr) | |
174 if len(stdout) > 4096: | |
175 logging.debug('Truncated output:') | |
176 logging.debug(stdout[:4096]) | |
177 return (status, stdout) | |
178 | |
179 def GetCmdStatusOutputAndError(args, cwd=None, shell=False): | |
180 """Executes a subprocess and returns its exit code, output, and errors. | |
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. Must be True if args | |
188 is a string and False if args is a sequence. | |
189 | |
190 Returns: | |
191 The 2-tuple (exit code, output). | |
192 """ | |
193 _ValidateAndLogCommand(args, cwd, shell) | |
194 pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, | |
195 shell=shell, cwd=cwd) | |
196 stdout, stderr = pipe.communicate() | |
197 return (pipe.returncode, stdout, stderr) | |
198 | |
199 | |
200 class TimeoutError(Exception): | |
201 """Module-specific timeout exception.""" | |
202 | |
203 def __init__(self, output=None): | |
204 super(TimeoutError, self).__init__() | |
205 self._output = output | |
206 | |
207 @property | |
208 def output(self): | |
209 return self._output | |
210 | |
211 | |
212 def _IterProcessStdout(process, timeout=None, buffer_size=4096, | |
213 poll_interval=1): | |
214 assert fcntl, 'fcntl module is required' | |
215 try: | |
216 # Enable non-blocking reads from the child's stdout. | |
217 child_fd = process.stdout.fileno() | |
218 fl = fcntl.fcntl(child_fd, fcntl.F_GETFL) | |
219 fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) | |
220 | |
221 end_time = (time.time() + timeout) if timeout else None | |
222 while True: | |
223 if end_time and time.time() > end_time: | |
224 raise TimeoutError() | |
225 read_fds, _, _ = select.select([child_fd], [], [], poll_interval) | |
226 if child_fd in read_fds: | |
227 data = os.read(child_fd, buffer_size) | |
228 if not data: | |
229 break | |
230 yield data | |
231 if process.poll() is not None: | |
232 break | |
233 finally: | |
234 try: | |
235 # Make sure the process doesn't stick around if we fail with an | |
236 # exception. | |
237 process.kill() | |
238 except OSError: | |
239 pass | |
240 process.wait() | |
241 | |
242 | |
243 def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False, | |
244 logfile=None): | |
245 """Executes a subprocess with a timeout. | |
246 | |
247 Args: | |
248 args: List of arguments to the program, the program to execute is the first | |
249 element. | |
250 timeout: the timeout in seconds or None to wait forever. | |
251 cwd: If not None, the subprocess's current directory will be changed to | |
252 |cwd| before it's executed. | |
253 shell: Whether to execute args as a shell command. Must be True if args | |
254 is a string and False if args is a sequence. | |
255 logfile: Optional file-like object that will receive output from the | |
256 command as it is running. | |
257 | |
258 Returns: | |
259 The 2-tuple (exit code, output). | |
260 """ | |
261 _ValidateAndLogCommand(args, cwd, shell) | |
262 output = StringIO.StringIO() | |
263 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, | |
264 stderr=subprocess.STDOUT) | |
265 try: | |
266 for data in _IterProcessStdout(process, timeout=timeout): | |
267 if logfile: | |
268 logfile.write(data) | |
269 output.write(data) | |
270 except TimeoutError: | |
271 raise TimeoutError(output.getvalue()) | |
272 | |
273 return process.returncode, output.getvalue() | |
274 | |
275 | |
276 def IterCmdOutputLines(args, timeout=None, cwd=None, shell=False, | |
277 check_status=True): | |
278 """Executes a subprocess and continuously yields lines from its output. | |
279 | |
280 Args: | |
281 args: List of arguments to the program, the program to execute is the first | |
282 element. | |
283 cwd: If not None, the subprocess's current directory will be changed to | |
284 |cwd| before it's executed. | |
285 shell: Whether to execute args as a shell command. Must be True if args | |
286 is a string and False if args is a sequence. | |
287 check_status: A boolean indicating whether to check the exit status of the | |
288 process after all output has been read. | |
289 | |
290 Yields: | |
291 The output of the subprocess, line by line. | |
292 | |
293 Raises: | |
294 CalledProcessError if check_status is True and the process exited with a | |
295 non-zero exit status. | |
296 """ | |
297 cmd = _ValidateAndLogCommand(args, cwd, shell) | |
298 process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE, | |
299 stderr=subprocess.STDOUT) | |
300 buffer_output = '' | |
301 for data in _IterProcessStdout(process, timeout=timeout): | |
302 buffer_output += data | |
303 has_incomplete_line = buffer_output[-1] not in '\r\n' | |
304 lines = buffer_output.splitlines() | |
305 buffer_output = lines.pop() if has_incomplete_line else '' | |
306 for line in lines: | |
307 yield line | |
308 if buffer_output: | |
309 yield buffer_output | |
310 if check_status and process.returncode: | |
311 raise subprocess.CalledProcessError(process.returncode, cmd) | |
OLD | NEW |