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 StringIO
13import subprocess
14import time
15
16# fcntl is not available on Windows.
17try:
18  import fcntl
19except ImportError:
20  fcntl = None
21
22
23def Popen(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
24  return subprocess.Popen(
25      args=args, cwd=cwd, stdout=stdout, stderr=stderr,
26      shell=shell, close_fds=True, env=env,
27      preexec_fn=lambda: signal.signal(signal.SIGPIPE, signal.SIG_DFL))
28
29
30def Call(args, stdout=None, stderr=None, shell=None, cwd=None, env=None):
31  pipe = Popen(args, stdout=stdout, stderr=stderr, shell=shell, cwd=cwd,
32               env=env)
33  pipe.communicate()
34  return pipe.wait()
35
36
37def RunCmd(args, cwd=None):
38  """Opens a subprocess to execute a program and returns its return value.
39
40  Args:
41    args: A string or a sequence of program arguments. The program to execute is
42      the string or the first item in the args sequence.
43    cwd: If not None, the subprocess's current directory will be changed to
44      |cwd| before it's executed.
45
46  Returns:
47    Return code from the command execution.
48  """
49  logging.info(str(args) + ' ' + (cwd or ''))
50  return Call(args, cwd=cwd)
51
52
53def GetCmdOutput(args, cwd=None, shell=False):
54  """Open a subprocess to execute a program and returns its output.
55
56  Args:
57    args: A string or a sequence of program arguments. The program to execute is
58      the string or the first item in the args sequence.
59    cwd: If not None, the subprocess's current directory will be changed to
60      |cwd| before it's executed.
61    shell: Whether to execute args as a shell command.
62
63  Returns:
64    Captures and returns the command's stdout.
65    Prints the command's stderr to logger (which defaults to stdout).
66  """
67  (_, output) = GetCmdStatusAndOutput(args, cwd, shell)
68  return output
69
70
71def GetCmdStatusAndOutput(args, cwd=None, shell=False):
72  """Executes a subprocess and returns its exit code and output.
73
74  Args:
75    args: A string or a sequence of program arguments. The program to execute is
76      the string or the first item in the args sequence.
77    cwd: If not None, the subprocess's current directory will be changed to
78      |cwd| before it's executed.
79    shell: Whether to execute args as a shell command.
80
81  Returns:
82    The 2-tuple (exit code, output).
83  """
84  if isinstance(args, basestring):
85    args_repr = args
86    if not shell:
87      raise Exception('string args must be run with shell=True')
88  elif shell:
89    raise Exception('array args must be run with shell=False')
90  else:
91    args_repr = ' '.join(map(pipes.quote, args))
92
93  s = '[host]'
94  if cwd:
95    s += ':' + cwd
96  s += '> ' + args_repr
97  logging.info(s)
98  pipe = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
99               shell=shell, cwd=cwd)
100  stdout, stderr = pipe.communicate()
101
102  if stderr:
103    logging.critical(stderr)
104  if len(stdout) > 4096:
105    logging.debug('Truncated output:')
106  logging.debug(stdout[:4096])
107  return (pipe.returncode, stdout)
108
109
110class TimeoutError(Exception):
111  """Module-specific timeout exception."""
112  pass
113
114
115def GetCmdStatusAndOutputWithTimeout(args, timeout, cwd=None, shell=False,
116                                     logfile=None):
117  """Executes a subprocess with a timeout.
118
119  Args:
120    args: List of arguments to the program, the program to execute is the first
121      element.
122    timeout: the timeout in seconds or None to wait forever.
123    cwd: If not None, the subprocess's current directory will be changed to
124      |cwd| before it's executed.
125    shell: Whether to execute args as a shell command.
126    logfile: Optional file-like object that will receive output from the
127      command as it is running.
128
129  Returns:
130    The 2-tuple (exit code, output).
131  """
132  assert fcntl, 'fcntl module is required'
133  process = Popen(args, cwd=cwd, shell=shell, stdout=subprocess.PIPE,
134                  stderr=subprocess.STDOUT)
135  try:
136    end_time = (time.time() + timeout) if timeout else None
137    poll_interval = 1
138    buffer_size = 4096
139    child_fd = process.stdout.fileno()
140    output = StringIO.StringIO()
141
142    # Enable non-blocking reads from the child's stdout.
143    fl = fcntl.fcntl(child_fd, fcntl.F_GETFL)
144    fcntl.fcntl(child_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
145
146    while True:
147      if end_time and time.time() > end_time:
148        raise TimeoutError
149      read_fds, _, _ = select.select([child_fd], [], [], poll_interval)
150      if child_fd in read_fds:
151        data = os.read(child_fd, buffer_size)
152        if not data:
153          break
154        if logfile:
155          logfile.write(data)
156        output.write(data)
157      if process.poll() is not None:
158        break
159  finally:
160    try:
161      # Make sure the process doesn't stick around if we fail with an
162      # exception.
163      process.kill()
164    except OSError:
165      pass
166    process.wait()
167  return process.returncode, output.getvalue()
168