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