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