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"""A wrapper for subprocess to make calling shell commands easier."""
5
6import codecs
7import logging
8import os
9import pipes
10import select
11import signal
12import string
13import subprocess
14import sys
15import time
16
17CATAPULT_ROOT_PATH = os.path.abspath(
18    os.path.join(os.path.dirname(__file__), '..', '..', '..'))
19SIX_PATH = os.path.join(CATAPULT_ROOT_PATH, 'third_party', 'six')
20
21if SIX_PATH not in sys.path:
22  sys.path.append(SIX_PATH)
23
24import six
25
26from devil import base_error
27
28logger = logging.getLogger(__name__)
29
30_SafeShellChars = frozenset(string.ascii_letters + string.digits + '@%_-+=:,./')
31
32# Cache the string-escape codec to ensure subprocess can find it
33# later. Return value doesn't matter.
34if six.PY2:
35  codecs.lookup('string-escape')
36
37
38def SingleQuote(s):
39  """Return an shell-escaped version of the string using single quotes.
40
41  Reliably quote a string which may contain unsafe characters (e.g. space,
42  quote, or other special characters such as '$').
43
44  The returned value can be used in a shell command line as one token that gets
45  to be interpreted literally.
46
47  Args:
48    s: The string to quote.
49
50  Return:
51    The string quoted using single quotes.
52  """
53  return pipes.quote(s)
54
55
56def DoubleQuote(s):
57  """Return an shell-escaped version of the string using double quotes.
58
59  Reliably quote a string which may contain unsafe characters (e.g. space
60  or quote characters), while retaining some shell features such as variable
61  interpolation.
62
63  The returned value can be used in a shell command line as one token that gets
64  to be further interpreted by the shell.
65
66  The set of characters that retain their special meaning may depend on the
67  shell implementation. This set usually includes: '$', '`', '\', '!', '*',
68  and '@'.
69
70  Args:
71    s: The string to quote.
72
73  Return:
74    The string quoted using double quotes.
75  """
76  if not s:
77    return '""'
78  elif all(c in _SafeShellChars for c in s):
79    return s
80  else:
81    return '"' + s.replace('"', '\\"') + '"'
82
83
84def ShrinkToSnippet(cmd_parts, var_name, var_value):
85  """Constructs a shell snippet for a command using a variable to shrink it.
86
87  Takes into account all quoting that needs to happen.
88
89  Args:
90    cmd_parts: A list of command arguments.
91    var_name: The variable that holds var_value.
92    var_value: The string to replace in cmd_parts with $var_name
93
94  Returns:
95    A shell snippet that does not include setting the variable.
96  """
97
98  def shrink(value):
99    parts = (x and SingleQuote(x) for x in value.split(var_value))
100    with_substitutions = ('"$%s"' % var_name).join(parts)
101    return with_substitutions or "''"
102
103  return ' '.join(shrink(part) for part in cmd_parts)
104
105
106def Popen(args,
107          stdin=None,
108          stdout=None,
109          stderr=None,
110          shell=None,
111          cwd=None,
112          env=None):
113  # preexec_fn isn't supported on windows.
114  # pylint: disable=unexpected-keyword-arg
115  if sys.platform == 'win32':
116    close_fds = (stdin is None and stdout is None and stderr is None)
117    preexec_fn = None
118  else:
119    close_fds = True
120    preexec_fn = lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL)
121
122  if six.PY2:
123    return subprocess.Popen(
124      args=args,
125      cwd=cwd,
126      stdin=stdin,
127      stdout=stdout,
128      stderr=stderr,
129      shell=shell,
130      close_fds=close_fds,
131      env=env,
132      preexec_fn=preexec_fn
133    )
134  else:
135    # opens stdout in text mode, so that caller side always get 'str',
136    # and there will be no type mismatch error.
137    # Ignore any decoding error, so that caller will not crash due to
138    # uncaught exception. Decoding errors are unavoidable, as we
139    # do not know the encoding of the output, and in some output there
140    # will be multiple encodings (e.g. adb logcat)
141    return subprocess.Popen(
142      args=args,
143      cwd=cwd,
144      stdin=stdin,
145      stdout=stdout,
146      stderr=stderr,
147      shell=shell,
148      close_fds=close_fds,
149      env=env,
150      preexec_fn=preexec_fn,
151      universal_newlines=True,
152      encoding='utf-8',
153      errors='ignore'
154    )
155
156def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
157  pipe = Popen(
158      args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd, env=env)
159  pipe.communicate()
160  return pipe.wait()
161
162
163def RunCmd(args, cwd=None):
164  """Opens a subprocess to execute a program and returns its return value.
165
166  Args:
167    args: A string or a sequence of program arguments. The program to execute is
168      the string or the first item in the args sequence.
169    cwd: If not None, the subprocess's current directory will be changed to
170      |cwd| before it's executed.
171
172  Returns:
173    Return code from the command execution.
174  """
175  logger.debug(str(args) + ' ' + (cwd or ''))
176  return Call(args, cwd=cwd)
177
178
179def GetCmdOutput(args, cwd=None, shell=False, env=None):
180  """Open a subprocess to execute a program and returns its output.
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.
188    env: If not None, a mapping that defines environment variables for the
189      subprocess.
190
191  Returns:
192    Captures and returns the command's stdout.
193    Prints the command's stderr to logger (which defaults to stdout).
194  """
195  (_, output) = GetCmdStatusAndOutput(args, cwd, shell, env)
196  return output
197
198
199def _ValidateAndLogCommand(args, cwd, shell):
200  if isinstance(args, six.string_types):
201    if not shell:
202      raise Exception('string args must be run with shell=True')
203  else:
204    if shell:
205      raise Exception('array args must be run with shell=False')
206    args = ' '.join(SingleQuote(str(c)) for c in args)
207  if cwd is None:
208    cwd = ''
209  else:
210    cwd = ':' + cwd
211  logger.debug('[host]%s> %s', cwd, args)
212  return args
213
214
215def GetCmdStatusAndOutput(args,
216                          cwd=None,
217                          shell=False,
218                          env=None,
219                          merge_stderr=False):
220  """Executes a subprocess and returns its exit code and output.
221
222  Args:
223    args: A string or a sequence of program arguments. The program to execute is
224      the string or the first item in the args sequence.
225    cwd: If not None, the subprocess's current directory will be changed to
226      |cwd| before it's executed.
227    shell: Whether to execute args as a shell command. Must be True if args
228      is a string and False if args is a sequence.
229    env: If not None, a mapping that defines environment variables for the
230      subprocess.
231    merge_stderr: If True, captures stderr as part of stdout.
232
233  Returns:
234    The 2-tuple (exit code, stdout).
235  """
236  status, stdout, stderr = GetCmdStatusOutputAndError(
237      args, cwd=cwd, shell=shell, env=env, merge_stderr=merge_stderr)
238
239  if stderr:
240    logger.critical('STDERR: %s', stderr)
241  logger.debug('STDOUT: %s%s', stdout[:4096].rstrip(),
242               '<truncated>' if len(stdout) > 4096 else '')
243  return (status, stdout)
244
245
246def StartCmd(args, cwd=None, shell=False, env=None):
247  """Starts a subprocess and returns a handle to the process.
248
249  Args:
250    args: A string or a sequence of program arguments. The program to execute is
251      the string or the first item in the args sequence.
252    cwd: If not None, the subprocess's current directory will be changed to
253      |cwd| before it's executed.
254    shell: Whether to execute args as a shell command. Must be True if args
255      is a string and False if args is a sequence.
256    env: If not None, a mapping that defines environment variables for the
257      subprocess.
258
259  Returns:
260    A process handle from subprocess.Popen.
261  """
262  _ValidateAndLogCommand(args, cwd, shell)
263  return Popen(
264      args,
265      stdout=subprocess.PIPE,
266      stderr=subprocess.PIPE,
267      shell=shell,
268      cwd=cwd,
269      env=env)
270
271
272def GetCmdStatusOutputAndError(args,
273                               cwd=None,
274                               shell=False,
275                               env=None,
276                               merge_stderr=False):
277  """Executes a subprocess and returns its exit code, output, and errors.
278
279  Args:
280    args: A string or a sequence of program arguments. The program to execute is
281      the string or the first item in the args sequence.
282    cwd: If not None, the subprocess's current directory will be changed to
283      |cwd| before it's executed.
284    shell: Whether to execute args as a shell command. Must be True if args
285      is a string and False if args is a sequence.
286    env: If not None, a mapping that defines environment variables for the
287      subprocess.
288    merge_stderr: If True, captures stderr as part of stdout.
289
290  Returns:
291    The 3-tuple (exit code, stdout, stderr).
292  """
293  _ValidateAndLogCommand(args, cwd, shell)
294  stderr = subprocess.STDOUT if merge_stderr else subprocess.PIPE
295  pipe = Popen(
296      args,
297      stdout=subprocess.PIPE,
298      stderr=stderr,
299      shell=shell,
300      cwd=cwd,
301      env=env)
302  stdout, stderr = pipe.communicate()
303  return (pipe.returncode, stdout, stderr)
304
305
306class TimeoutError(base_error.BaseError):
307  """Module-specific timeout exception."""
308
309  def __init__(self, output=None):
310    super(TimeoutError, self).__init__('Timeout')
311    self._output = output
312
313  @property
314  def output(self):
315    return self._output
316
317
318def _read_and_decode(fd, buffer_size):
319  data = os.read(fd, buffer_size)
320  if data and six.PY3:
321    data = data.decode('utf-8', errors='ignore')
322  return data
323
324def _IterProcessStdoutFcntl(process,
325                            iter_timeout=None,
326                            timeout=None,
327                            buffer_size=4096,
328                            poll_interval=1):
329  """An fcntl-based implementation of _IterProcessStdout."""
330  # pylint: disable=too-many-nested-blocks
331  import fcntl
332  try:
333    # Enable non-blocking reads from the child's stdout.
334    child_fd = process.stdout.fileno()
335    fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
336    fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
337
338    end_time = (time.time() + timeout) if timeout else None
339    iter_end_time = (time.time() + iter_timeout) if iter_timeout else None
340
341    while True:
342      if end_time and time.time() > end_time:
343        raise TimeoutError()
344      if iter_end_time and time.time() > iter_end_time:
345        yield None
346        iter_end_time = time.time() + iter_timeout
347
348      if iter_end_time:
349        iter_aware_poll_interval = min(poll_interval,
350                                       max(0, iter_end_time - time.time()))
351      else:
352        iter_aware_poll_interval = poll_interval
353
354      read_fds, _, _ = select.select([child_fd], [], [],
355                                     iter_aware_poll_interval)
356      if child_fd in read_fds:
357        data = _read_and_decode(child_fd, buffer_size)
358        if not data:
359          break
360        yield data
361
362      if process.poll() is not None:
363        # If process is closed, keep checking for output data (because of timing
364        # issues).
365        while True:
366          read_fds, _, _ = select.select([child_fd], [], [],
367                                         iter_aware_poll_interval)
368          if child_fd in read_fds:
369            data = _read_and_decode(child_fd, buffer_size)
370            if data:
371              yield data
372              continue
373          break
374        break
375  finally:
376    try:
377      if process.returncode is None:
378        # Make sure the process doesn't stick around if we fail with an
379        # exception.
380        process.kill()
381    except OSError:
382      pass
383    process.wait()
384
385
386def _IterProcessStdoutQueue(process,
387                            iter_timeout=None,
388                            timeout=None,
389                            buffer_size=4096,
390                            poll_interval=1):
391  """A Queue.Queue-based implementation of _IterProcessStdout.
392
393  TODO(jbudorick): Evaluate whether this is a suitable replacement for
394  _IterProcessStdoutFcntl on all platforms.
395  """
396  # pylint: disable=unused-argument
397  import Queue
398  import threading
399
400  stdout_queue = Queue.Queue()
401
402  def read_process_stdout():
403    # TODO(jbudorick): Pick an appropriate read size here.
404    while True:
405      try:
406        output_chunk = _read_and_decode(process.stdout.fileno(), buffer_size)
407      except IOError:
408        break
409      stdout_queue.put(output_chunk, True)
410      if not output_chunk and process.poll() is not None:
411        break
412
413  reader_thread = threading.Thread(target=read_process_stdout)
414  reader_thread.start()
415
416  end_time = (time.time() + timeout) if timeout else None
417
418  try:
419    while True:
420      if end_time and time.time() > end_time:
421        raise TimeoutError()
422      try:
423        s = stdout_queue.get(True, iter_timeout)
424        if not s:
425          break
426        yield s
427      except Queue.Empty:
428        yield None
429  finally:
430    try:
431      if process.returncode is None:
432        # Make sure the process doesn't stick around if we fail with an
433        # exception.
434        process.kill()
435    except OSError:
436      pass
437    process.wait()
438    reader_thread.join()
439
440
441_IterProcessStdout = (_IterProcessStdoutQueue
442                      if sys.platform == 'win32' else _IterProcessStdoutFcntl)
443"""Iterate over a process's stdout.
444
445This is intentionally not public.
446
447Args:
448  process: The process in question.
449  iter_timeout: An optional length of time, in seconds, to wait in
450    between each iteration. If no output is received in the given
451    time, this generator will yield None.
452  timeout: An optional length of time, in seconds, during which
453    the process must finish. If it fails to do so, a TimeoutError
454    will be raised.
455  buffer_size: The maximum number of bytes to read (and thus yield) at once.
456  poll_interval: The length of time to wait in calls to `select.select`.
457    If iter_timeout is set, the remaining length of time in the iteration
458    may take precedence.
459Raises:
460  TimeoutError: if timeout is set and the process does not complete.
461Yields:
462  basestrings of data or None.
463"""
464
465
466def GetCmdStatusAndOutputWithTimeout(args,
467                                     timeout,
468                                     cwd=None,
469                                     shell=False,
470                                     logfile=None,
471                                     env=None):
472  """Executes a subprocess with a timeout.
473
474  Args:
475    args: List of arguments to the program, the program to execute is the first
476      element.
477    timeout: the timeout in seconds or None to wait forever.
478    cwd: If not None, the subprocess's current directory will be changed to
479      |cwd| before it's executed.
480    shell: Whether to execute args as a shell command. Must be True if args
481      is a string and False if args is a sequence.
482    logfile: Optional file-like object that will receive output from the
483      command as it is running.
484    env: If not None, a mapping that defines environment variables for the
485      subprocess.
486
487  Returns:
488    The 2-tuple (exit code, output).
489  Raises:
490    TimeoutError on timeout.
491  """
492  _ValidateAndLogCommand(args, cwd, shell)
493  output = six.StringIO()
494  process = Popen(
495      args,
496      cwd=cwd,
497      shell=shell,
498      stdout=subprocess.PIPE,
499      stderr=subprocess.STDOUT,
500      env=env)
501  try:
502    for data in _IterProcessStdout(process, timeout=timeout):
503      if logfile:
504        logfile.write(data)
505      output.write(data)
506  except TimeoutError:
507    raise TimeoutError(output.getvalue())
508
509  str_output = output.getvalue()
510  logger.debug('STDOUT+STDERR: %s%s', str_output[:4096].rstrip(),
511               '<truncated>' if len(str_output) > 4096 else '')
512  return process.returncode, str_output
513
514
515def IterCmdOutputLines(args,
516                       iter_timeout=None,
517                       timeout=None,
518                       cwd=None,
519                       shell=False,
520                       env=None,
521                       check_status=True):
522  """Executes a subprocess and continuously yields lines from its output.
523
524  Args:
525    args: List of arguments to the program, the program to execute is the first
526      element.
527    iter_timeout: Timeout for each iteration, in seconds.
528    timeout: Timeout for the entire command, in seconds.
529    cwd: If not None, the subprocess's current directory will be changed to
530      |cwd| before it's executed.
531    shell: Whether to execute args as a shell command. Must be True if args
532      is a string and False if args is a sequence.
533    env: If not None, a mapping that defines environment variables for the
534      subprocess.
535    check_status: A boolean indicating whether to check the exit status of the
536      process after all output has been read.
537  Yields:
538    The output of the subprocess, line by line.
539
540  Raises:
541    CalledProcessError if check_status is True and the process exited with a
542      non-zero exit status.
543  """
544  cmd = _ValidateAndLogCommand(args, cwd, shell)
545  process = Popen(
546      args,
547      cwd=cwd,
548      shell=shell,
549      env=env,
550      stdout=subprocess.PIPE,
551      stderr=subprocess.STDOUT)
552  return _IterCmdOutputLines(
553      process,
554      cmd,
555      iter_timeout=iter_timeout,
556      timeout=timeout,
557      check_status=check_status)
558
559
560def _IterCmdOutputLines(process,
561                        cmd,
562                        iter_timeout=None,
563                        timeout=None,
564                        check_status=True):
565  buffer_output = ''
566
567  iter_end = None
568  cur_iter_timeout = None
569  if iter_timeout:
570    iter_end = time.time() + iter_timeout
571    cur_iter_timeout = iter_timeout
572
573  for data in _IterProcessStdout(
574      process, iter_timeout=cur_iter_timeout, timeout=timeout):
575    if iter_timeout:
576      # Check whether the current iteration has timed out.
577      cur_iter_timeout = iter_end - time.time()
578      if data is None or cur_iter_timeout < 0:
579        yield None
580        iter_end = time.time() + iter_timeout
581        continue
582    else:
583      assert data is not None, (
584          'Iteration received no data despite no iter_timeout being set. '
585          'cmd: %s' % cmd)
586
587    # Construct lines to yield from raw data.
588    buffer_output += data
589    has_incomplete_line = buffer_output[-1] not in '\r\n'
590    lines = buffer_output.splitlines()
591    buffer_output = lines.pop() if has_incomplete_line else ''
592    for line in lines:
593      yield line
594      if iter_timeout:
595        iter_end = time.time() + iter_timeout
596
597  if buffer_output:
598    yield buffer_output
599  if check_status and process.returncode:
600    raise subprocess.CalledProcessError(process.returncode, cmd)
601