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: %s', stderr)
176  logging.debug('STDOUT: %s%s', stdout[:4096].rstrip(),
177                '<truncated>' if len(stdout) > 4096 else '')
178  return (status, stdout)
179
180
181def GetCmdStatusOutputAndError(args, cwd=None, shell=False):
182  """Executes a subprocess and returns its exit code, output, and errors.
183
184  Args:
185    args: A string or a sequence of program arguments. The program to execute is
186      the string or the first item in the args sequence.
187    cwd: If not None, the subprocess's current directory will be changed to
188      |cwd| before it's executed.
189    shell: Whether to execute args as a shell command. Must be True if args
190      is a string and False if args is a sequence.
191
192  Returns:
193    The 2-tuple (exit code, output).
194  """
195  _ValidateAndLogCommand(args, cwd, shell)
196  pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
197               shell=shell, cwd=cwd)
198  stdout, stderr = pipe.communicate()
199  return (pipe.returncode, stdout, stderr)
200
201
202class TimeoutError(Exception):
203  """Module-specific timeout exception."""
204
205  def __init__(self, output=None):
206    super(TimeoutError, self).__init__()
207    self._output = output
208
209  @property
210  def output(self):
211    return self._output
212
213
214def _IterProcessStdout(process, timeout=None, buffer_size=4096,
215                       poll_interval=1):
216  assert fcntl, 'fcntl module is required'
217  try:
218    # Enable non-blocking reads from the child's stdout.
219    child_fd = process.stdout.fileno()
220    fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
221    fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
222
223    end_time = (time.time() + timeout) if timeout else None
224    while True:
225      if end_time and time.time() > end_time:
226        raise TimeoutError()
227      read_fds, _, _ = select.select([child_fd], [], [], poll_interval)
228      if child_fd in read_fds:
229        data = os.read(child_fd, buffer_size)
230        if not data:
231          break
232        yield data
233      if process.poll() is not None:
234        break
235  finally:
236    try:
237      # Make sure the process doesn't stick around if we fail with an
238      # exception.
239      process.kill()
240    except OSError:
241      pass
242    process.wait()
243
244
245def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False,
246                                     logfile=None):
247  """Executes a subprocess with a timeout.
248
249  Args:
250    args: List of arguments to the program, the program to execute is the first
251      element.
252    timeout: the timeout in seconds or None to wait forever.
253    cwd: If not None, the subprocess's current directory will be changed to
254      |cwd| before it's executed.
255    shell: Whether to execute args as a shell command. Must be True if args
256      is a string and False if args is a sequence.
257    logfile: Optional file-like object that will receive output from the
258      command as it is running.
259
260  Returns:
261    The 2-tuple (exit code, output).
262  """
263  _ValidateAndLogCommand(args, cwd, shell)
264  output = StringIO.StringIO()
265  process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
266                  stderr=subprocess.STDOUT)
267  try:
268    for data in _IterProcessStdout(process, timeout=timeout):
269      if logfile:
270        logfile.write(data)
271      output.write(data)
272  except TimeoutError:
273    raise TimeoutError(output.getvalue())
274
275  str_output = output.getvalue()
276  logging.debug('STDOUT+STDERR: %s%s', str_output[:4096].rstrip(),
277                '<truncated>' if len(str_output) > 4096 else '')
278  return process.returncode, str_output
279
280
281def IterCmdOutputLines(args, timeout=None, cwd=None, shell=False,
282                       check_status=True):
283  """Executes a subprocess and continuously yields lines from its output.
284
285  Args:
286    args: List of arguments to the program, the program to execute is the first
287      element.
288    cwd: If not None, the subprocess's current directory will be changed to
289      |cwd| before it's executed.
290    shell: Whether to execute args as a shell command. Must be True if args
291      is a string and False if args is a sequence.
292    check_status: A boolean indicating whether to check the exit status of the
293      process after all output has been read.
294
295  Yields:
296    The output of the subprocess, line by line.
297
298  Raises:
299    CalledProcessError if check_status is True and the process exited with a
300      non-zero exit status.
301  """
302  cmd = _ValidateAndLogCommand(args, cwd, shell)
303  process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
304                  stderr=subprocess.STDOUT)
305  buffer_output = ''
306  for data in _IterProcessStdout(process, timeout=timeout):
307    buffer_output += data
308    has_incomplete_line = buffer_output[-1] not in '\r\n'
309    lines = buffer_output.splitlines()
310    buffer_output = lines.pop() if has_incomplete_line else ''
311    for line in lines:
312      yield line
313  if buffer_output:
314    yield buffer_output
315  if check_status and process.returncode:
316    raise subprocess.CalledProcessError(process.returncode, cmd)
317