1# Copyright 2016 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Various utility functions.""" 16 17import errno 18import functools 19import os 20import signal 21import subprocess 22import sys 23import tempfile 24import time 25 26_path = os.path.realpath(__file__ + '/../..') 27if sys.path[0] != _path: 28 sys.path.insert(0, _path) 29del _path 30 31# pylint: disable=wrong-import-position 32import rh.shell 33import rh.signals 34 35 36def timedelta_str(delta): 37 """A less noisy timedelta.__str__. 38 39 The default timedelta stringification contains a lot of leading zeros and 40 uses microsecond resolution. This makes for noisy output. 41 """ 42 total = delta.total_seconds() 43 hours, rem = divmod(total, 3600) 44 mins, secs = divmod(rem, 60) 45 ret = '%i.%03is' % (secs, delta.microseconds // 1000) 46 if mins: 47 ret = '%im%s' % (mins, ret) 48 if hours: 49 ret = '%ih%s' % (hours, ret) 50 return ret 51 52 53class CompletedProcess(getattr(subprocess, 'CompletedProcess', object)): 54 """An object to store various attributes of a child process. 55 56 This is akin to subprocess.CompletedProcess. 57 """ 58 59 # The linter is confused by the getattr usage above. 60 # TODO(vapier): Drop this once we're Python 3-only and we drop getattr. 61 # pylint: disable=bad-option-value,super-on-old-class 62 def __init__(self, args=None, returncode=None, stdout=None, stderr=None): 63 if sys.version_info.major < 3: 64 self.args = args 65 self.stdout = stdout 66 self.stderr = stderr 67 self.returncode = returncode 68 else: 69 super(CompletedProcess, self).__init__( 70 args=args, returncode=returncode, stdout=stdout, stderr=stderr) 71 72 @property 73 def cmd(self): 74 """Alias to self.args to better match other subprocess APIs.""" 75 return self.args 76 77 @property 78 def cmdstr(self): 79 """Return self.cmd as a nicely formatted string (useful for logs).""" 80 return rh.shell.cmd_to_str(self.cmd) 81 82 83class CalledProcessError(subprocess.CalledProcessError): 84 """Error caught in run() function. 85 86 This is akin to subprocess.CalledProcessError. We do not support |output|, 87 only |stdout|. 88 89 Attributes: 90 returncode: The exit code of the process. 91 cmd: The command that triggered this exception. 92 msg: Short explanation of the error. 93 exception: The underlying Exception if available. 94 """ 95 96 def __init__(self, returncode, cmd, stdout=None, stderr=None, msg=None, 97 exception=None): 98 if exception is not None and not isinstance(exception, Exception): 99 raise TypeError('exception must be an exception instance; got %r' 100 % (exception,)) 101 102 super(CalledProcessError, self).__init__(returncode, cmd, stdout) 103 # The parent class will set |output|, so delete it. 104 del self.output 105 # TODO(vapier): When we're Python 3-only, delete this assignment as the 106 # parent handles it for us. 107 self.stdout = stdout 108 # TODO(vapier): When we're Python 3-only, move stderr to the init above. 109 self.stderr = stderr 110 self.msg = msg 111 self.exception = exception 112 113 @property 114 def cmdstr(self): 115 """Return self.cmd as a well shell-quoted string for debugging.""" 116 return '' if self.cmd is None else rh.shell.cmd_to_str(self.cmd) 117 118 def stringify(self, stdout=True, stderr=True): 119 """Custom method for controlling what is included in stringifying this. 120 121 Args: 122 stdout: Whether to include captured stdout in the return value. 123 stderr: Whether to include captured stderr in the return value. 124 125 Returns: 126 A summary string for this result. 127 """ 128 items = [ 129 'return code: %s; command: %s' % (self.returncode, self.cmdstr), 130 ] 131 if stderr and self.stderr: 132 items.append(self.stderr) 133 if stdout and self.stdout: 134 items.append(self.stdout) 135 if self.msg: 136 items.append(self.msg) 137 return '\n'.join(items) 138 139 def __str__(self): 140 return self.stringify() 141 142 143class TerminateCalledProcessError(CalledProcessError): 144 """We were signaled to shutdown while running a command. 145 146 Client code shouldn't generally know, nor care about this class. It's 147 used internally to suppress retry attempts when we're signaled to die. 148 """ 149 150 151def _kill_child_process(proc, int_timeout, kill_timeout, cmd, original_handler, 152 signum, frame): 153 """Used as a signal handler by RunCommand. 154 155 This is internal to Runcommand. No other code should use this. 156 """ 157 if signum: 158 # If we've been invoked because of a signal, ignore delivery of that 159 # signal from this point forward. The invoking context of this func 160 # restores signal delivery to what it was prior; we suppress future 161 # delivery till then since this code handles SIGINT/SIGTERM fully 162 # including delivering the signal to the original handler on the way 163 # out. 164 signal.signal(signum, signal.SIG_IGN) 165 166 # Do not trust Popen's returncode alone; we can be invoked from contexts 167 # where the Popen instance was created, but no process was generated. 168 if proc.returncode is None and proc.pid is not None: 169 try: 170 while proc.poll_lock_breaker() is None and int_timeout >= 0: 171 time.sleep(0.1) 172 int_timeout -= 0.1 173 174 proc.terminate() 175 while proc.poll_lock_breaker() is None and kill_timeout >= 0: 176 time.sleep(0.1) 177 kill_timeout -= 0.1 178 179 if proc.poll_lock_breaker() is None: 180 # Still doesn't want to die. Too bad, so sad, time to die. 181 proc.kill() 182 except EnvironmentError as e: 183 print('Ignoring unhandled exception in _kill_child_process: %s' % e, 184 file=sys.stderr) 185 186 # Ensure our child process has been reaped. 187 kwargs = {} 188 if sys.version_info.major >= 3: 189 # ... but don't wait forever. 190 kwargs['timeout'] = 60 191 proc.wait_lock_breaker(**kwargs) 192 193 if not rh.signals.relay_signal(original_handler, signum, frame): 194 # Mock up our own, matching exit code for signaling. 195 raise TerminateCalledProcessError( 196 signum << 8, cmd, msg='Received signal %i' % signum) 197 198 199class _Popen(subprocess.Popen): 200 """subprocess.Popen derivative customized for our usage. 201 202 Specifically, we fix terminate/send_signal/kill to work if the child process 203 was a setuid binary; on vanilla kernels, the parent can wax the child 204 regardless, on goobuntu this apparently isn't allowed, thus we fall back 205 to the sudo machinery we have. 206 207 While we're overriding send_signal, we also suppress ESRCH being raised 208 if the process has exited, and suppress signaling all together if the 209 process has knowingly been waitpid'd already. 210 """ 211 212 # pylint: disable=arguments-differ 213 def send_signal(self, signum): 214 if self.returncode is not None: 215 # The original implementation in Popen allows signaling whatever 216 # process now occupies this pid, even if the Popen object had 217 # waitpid'd. Since we can escalate to sudo kill, we do not want 218 # to allow that. Fixing this addresses that angle, and makes the 219 # API less sucky in the process. 220 return 221 222 try: 223 os.kill(self.pid, signum) 224 except EnvironmentError as e: 225 if e.errno == errno.ESRCH: 226 # Since we know the process is dead, reap it now. 227 # Normally Popen would throw this error- we suppress it since 228 # frankly that's a misfeature and we're already overriding 229 # this method. 230 self.poll() 231 else: 232 raise 233 234 def _lock_breaker(self, func, *args, **kwargs): 235 """Helper to manage the waitpid lock. 236 237 Workaround https://bugs.python.org/issue25960. 238 """ 239 # If the lock doesn't exist, or is not locked, call the func directly. 240 lock = getattr(self, '_waitpid_lock', None) 241 if lock is not None and lock.locked(): 242 try: 243 lock.release() 244 return func(*args, **kwargs) 245 finally: 246 if not lock.locked(): 247 lock.acquire() 248 else: 249 return func(*args, **kwargs) 250 251 def poll_lock_breaker(self, *args, **kwargs): 252 """Wrapper around poll() to break locks if needed.""" 253 return self._lock_breaker(self.poll, *args, **kwargs) 254 255 def wait_lock_breaker(self, *args, **kwargs): 256 """Wrapper around wait() to break locks if needed.""" 257 return self._lock_breaker(self.wait, *args, **kwargs) 258 259 260# We use the keyword arg |input| which trips up pylint checks. 261# pylint: disable=redefined-builtin,input-builtin 262def run(cmd, redirect_stdout=False, redirect_stderr=False, cwd=None, input=None, 263 shell=False, env=None, extra_env=None, combine_stdout_stderr=False, 264 check=True, int_timeout=1, kill_timeout=1, capture_output=False, 265 close_fds=True): 266 """Runs a command. 267 268 Args: 269 cmd: cmd to run. Should be input to subprocess.Popen. If a string, shell 270 must be true. Otherwise the command must be an array of arguments, 271 and shell must be false. 272 redirect_stdout: Returns the stdout. 273 redirect_stderr: Holds stderr output until input is communicated. 274 cwd: The working directory to run this cmd. 275 input: The data to pipe into this command through stdin. If a file object 276 or file descriptor, stdin will be connected directly to that. 277 shell: Controls whether we add a shell as a command interpreter. See cmd 278 since it has to agree as to the type. 279 env: If non-None, this is the environment for the new process. 280 extra_env: If set, this is added to the environment for the new process. 281 This dictionary is not used to clear any entries though. 282 combine_stdout_stderr: Combines stdout and stderr streams into stdout. 283 check: Whether to raise an exception when command returns a non-zero exit 284 code, or return the CompletedProcess object containing the exit code. 285 Note: will still raise an exception if the cmd file does not exist. 286 int_timeout: If we're interrupted, how long (in seconds) should we give 287 the invoked process to clean up before we send a SIGTERM. 288 kill_timeout: If we're interrupted, how long (in seconds) should we give 289 the invoked process to shutdown from a SIGTERM before we SIGKILL it. 290 capture_output: Set |redirect_stdout| and |redirect_stderr| to True. 291 close_fds: Whether to close all fds before running |cmd|. 292 293 Returns: 294 A CompletedProcess object. 295 296 Raises: 297 CalledProcessError: Raises exception on error. 298 """ 299 if capture_output: 300 redirect_stdout, redirect_stderr = True, True 301 302 # Set default for variables. 303 popen_stdout = None 304 popen_stderr = None 305 stdin = None 306 result = CompletedProcess() 307 308 # Force the timeout to float; in the process, if it's not convertible, 309 # a self-explanatory exception will be thrown. 310 kill_timeout = float(kill_timeout) 311 312 def _get_tempfile(): 313 kwargs = {} 314 if sys.version_info.major < 3: 315 kwargs['bufsize'] = 0 316 else: 317 kwargs['buffering'] = 0 318 try: 319 return tempfile.TemporaryFile(**kwargs) 320 except EnvironmentError as e: 321 if e.errno != errno.ENOENT: 322 raise 323 # This can occur if we were pointed at a specific location for our 324 # TMP, but that location has since been deleted. Suppress that 325 # issue in this particular case since our usage gurantees deletion, 326 # and since this is primarily triggered during hard cgroups 327 # shutdown. 328 return tempfile.TemporaryFile(dir='/tmp', **kwargs) 329 330 # Modify defaults based on parameters. 331 # Note that tempfiles must be unbuffered else attempts to read 332 # what a separate process did to that file can result in a bad 333 # view of the file. 334 # The Popen API accepts either an int or a file handle for stdout/stderr. 335 # pylint: disable=redefined-variable-type 336 if redirect_stdout: 337 popen_stdout = _get_tempfile() 338 339 if combine_stdout_stderr: 340 popen_stderr = subprocess.STDOUT 341 elif redirect_stderr: 342 popen_stderr = _get_tempfile() 343 # pylint: enable=redefined-variable-type 344 345 # If subprocesses have direct access to stdout or stderr, they can bypass 346 # our buffers, so we need to flush to ensure that output is not interleaved. 347 if popen_stdout is None or popen_stderr is None: 348 sys.stdout.flush() 349 sys.stderr.flush() 350 351 # If input is a string, we'll create a pipe and send it through that. 352 # Otherwise we assume it's a file object that can be read from directly. 353 if isinstance(input, str): 354 stdin = subprocess.PIPE 355 input = input.encode('utf-8') 356 elif input is not None: 357 stdin = input 358 input = None 359 360 if isinstance(cmd, str): 361 if not shell: 362 raise Exception('Cannot run a string command without a shell') 363 cmd = ['/bin/bash', '-c', cmd] 364 shell = False 365 elif shell: 366 raise Exception('Cannot run an array command with a shell') 367 368 # If we are using enter_chroot we need to use enterchroot pass env through 369 # to the final command. 370 env = env.copy() if env is not None else os.environ.copy() 371 env.update(extra_env if extra_env else {}) 372 373 result.args = cmd 374 375 proc = None 376 try: 377 proc = _Popen(cmd, cwd=cwd, stdin=stdin, stdout=popen_stdout, 378 stderr=popen_stderr, shell=False, env=env, 379 close_fds=close_fds) 380 381 old_sigint = signal.getsignal(signal.SIGINT) 382 handler = functools.partial(_kill_child_process, proc, int_timeout, 383 kill_timeout, cmd, old_sigint) 384 signal.signal(signal.SIGINT, handler) 385 386 old_sigterm = signal.getsignal(signal.SIGTERM) 387 handler = functools.partial(_kill_child_process, proc, int_timeout, 388 kill_timeout, cmd, old_sigterm) 389 signal.signal(signal.SIGTERM, handler) 390 391 try: 392 (result.stdout, result.stderr) = proc.communicate(input) 393 finally: 394 signal.signal(signal.SIGINT, old_sigint) 395 signal.signal(signal.SIGTERM, old_sigterm) 396 397 if popen_stdout: 398 # The linter is confused by how stdout is a file & an int. 399 # pylint: disable=maybe-no-member,no-member 400 popen_stdout.seek(0) 401 result.stdout = popen_stdout.read() 402 popen_stdout.close() 403 404 if popen_stderr and popen_stderr != subprocess.STDOUT: 405 # The linter is confused by how stderr is a file & an int. 406 # pylint: disable=maybe-no-member,no-member 407 popen_stderr.seek(0) 408 result.stderr = popen_stderr.read() 409 popen_stderr.close() 410 411 result.returncode = proc.returncode 412 413 if check and proc.returncode: 414 msg = 'cwd=%s' % cwd 415 if extra_env: 416 msg += ', extra env=%s' % extra_env 417 raise CalledProcessError( 418 result.returncode, result.cmd, stdout=result.stdout, 419 stderr=result.stderr, msg=msg) 420 except OSError as e: 421 estr = str(e) 422 if e.errno == errno.EACCES: 423 estr += '; does the program need `chmod a+x`?' 424 if not check: 425 result = CompletedProcess( 426 args=cmd, stderr=estr.encode('utf-8'), returncode=255) 427 else: 428 raise CalledProcessError( 429 result.returncode, result.cmd, stdout=result.stdout, 430 stderr=result.stderr, msg=estr, exception=e) 431 finally: 432 if proc is not None: 433 # Ensure the process is dead. 434 # Some pylint3 versions are confused here. 435 # pylint: disable=too-many-function-args 436 _kill_child_process(proc, int_timeout, kill_timeout, cmd, None, 437 None, None) 438 439 # Make sure output is returned as a string rather than bytes. 440 if result.stdout is not None: 441 result.stdout = result.stdout.decode('utf-8', 'replace') 442 if result.stderr is not None: 443 result.stderr = result.stderr.decode('utf-8', 'replace') 444 445 return result 446# pylint: enable=redefined-builtin,input-builtin 447