1*9c5db199SXin Li# Lint as: python2, python3 2*9c5db199SXin Li# Copyright (c) 2013 The Chromium OS Authors. All rights reserved. 3*9c5db199SXin Li# Use of this source code is governed by a BSD-style license that can be 4*9c5db199SXin Li# found in the LICENSE file. 5*9c5db199SXin Li 6*9c5db199SXin Liimport atexit 7*9c5db199SXin Liimport itertools 8*9c5db199SXin Liimport logging 9*9c5db199SXin Liimport os 10*9c5db199SXin Liimport pipes 11*9c5db199SXin Liimport pwd 12*9c5db199SXin Liimport select 13*9c5db199SXin Liimport subprocess 14*9c5db199SXin Liimport threading 15*9c5db199SXin Li 16*9c5db199SXin Lifrom autotest_lib.client.common_lib.utils import TEE_TO_LOGS 17*9c5db199SXin Li 18*9c5db199SXin Li_popen_lock = threading.Lock() 19*9c5db199SXin Li_logging_service = None 20*9c5db199SXin Li_command_serial_number = itertools.count(1) 21*9c5db199SXin Li 22*9c5db199SXin Li_LOG_BUFSIZE = 4096 23*9c5db199SXin Li_PIPE_CLOSED = -1 24*9c5db199SXin Li 25*9c5db199SXin Liclass _LoggerProxy(object): 26*9c5db199SXin Li 27*9c5db199SXin Li def __init__(self, logger): 28*9c5db199SXin Li self._logger = logger 29*9c5db199SXin Li 30*9c5db199SXin Li def fileno(self): 31*9c5db199SXin Li """Returns the fileno of the logger pipe.""" 32*9c5db199SXin Li return self._logger._pipe[1] 33*9c5db199SXin Li 34*9c5db199SXin Li def __del__(self): 35*9c5db199SXin Li self._logger.close() 36*9c5db199SXin Li 37*9c5db199SXin Li 38*9c5db199SXin Liclass _PipeLogger(object): 39*9c5db199SXin Li 40*9c5db199SXin Li def __init__(self, level, prefix): 41*9c5db199SXin Li self._pipe = list(os.pipe()) 42*9c5db199SXin Li self._level = level 43*9c5db199SXin Li self._prefix = prefix 44*9c5db199SXin Li 45*9c5db199SXin Li def close(self): 46*9c5db199SXin Li """Closes the logger.""" 47*9c5db199SXin Li if self._pipe[1] != _PIPE_CLOSED: 48*9c5db199SXin Li os.close(self._pipe[1]) 49*9c5db199SXin Li self._pipe[1] = _PIPE_CLOSED 50*9c5db199SXin Li 51*9c5db199SXin Li 52*9c5db199SXin Liclass _LoggingService(object): 53*9c5db199SXin Li 54*9c5db199SXin Li def __init__(self): 55*9c5db199SXin Li # Python's list is thread safe 56*9c5db199SXin Li self._loggers = [] 57*9c5db199SXin Li 58*9c5db199SXin Li # Change tuple to list so that we can change the value when 59*9c5db199SXin Li # closing the pipe. 60*9c5db199SXin Li self._pipe = list(os.pipe()) 61*9c5db199SXin Li self._thread = threading.Thread(target=self._service_run) 62*9c5db199SXin Li self._thread.daemon = True 63*9c5db199SXin Li self._thread.start() 64*9c5db199SXin Li 65*9c5db199SXin Li 66*9c5db199SXin Li def _service_run(self): 67*9c5db199SXin Li terminate_loop = False 68*9c5db199SXin Li while not terminate_loop: 69*9c5db199SXin Li rlist = [l._pipe[0] for l in self._loggers] 70*9c5db199SXin Li rlist.append(self._pipe[0]) 71*9c5db199SXin Li for r in select.select(rlist, [], [])[0]: 72*9c5db199SXin Li data = os.read(r, _LOG_BUFSIZE) 73*9c5db199SXin Li if r != self._pipe[0]: 74*9c5db199SXin Li self._output_logger_message(r, data) 75*9c5db199SXin Li elif len(data) == 0: 76*9c5db199SXin Li terminate_loop = True 77*9c5db199SXin Li # Release resources. 78*9c5db199SXin Li os.close(self._pipe[0]) 79*9c5db199SXin Li for logger in self._loggers: 80*9c5db199SXin Li os.close(logger._pipe[0]) 81*9c5db199SXin Li 82*9c5db199SXin Li 83*9c5db199SXin Li def _output_logger_message(self, r, data): 84*9c5db199SXin Li logger = next(l for l in self._loggers if l._pipe[0] == r) 85*9c5db199SXin Li 86*9c5db199SXin Li if len(data) == 0: 87*9c5db199SXin Li os.close(logger._pipe[0]) 88*9c5db199SXin Li self._loggers.remove(logger) 89*9c5db199SXin Li return 90*9c5db199SXin Li 91*9c5db199SXin Li for line in data.split(b'\n'): 92*9c5db199SXin Li logging.log(logger._level, '%s%s', logger._prefix, line) 93*9c5db199SXin Li 94*9c5db199SXin Li 95*9c5db199SXin Li def create_logger(self, level=logging.DEBUG, prefix=''): 96*9c5db199SXin Li """Creates a new logger. 97*9c5db199SXin Li 98*9c5db199SXin Li @param level: the desired logging level 99*9c5db199SXin Li @param prefix: the prefix to add to each log entry 100*9c5db199SXin Li """ 101*9c5db199SXin Li logger = _PipeLogger(level=level, prefix=prefix) 102*9c5db199SXin Li self._loggers.append(logger) 103*9c5db199SXin Li os.write(self._pipe[1], b'\0') 104*9c5db199SXin Li return _LoggerProxy(logger) 105*9c5db199SXin Li 106*9c5db199SXin Li 107*9c5db199SXin Li def shutdown(self): 108*9c5db199SXin Li """Shuts down the logger.""" 109*9c5db199SXin Li if self._pipe[1] != _PIPE_CLOSED: 110*9c5db199SXin Li os.close(self._pipe[1]) 111*9c5db199SXin Li self._pipe[1] = _PIPE_CLOSED 112*9c5db199SXin Li self._thread.join() 113*9c5db199SXin Li 114*9c5db199SXin Li 115*9c5db199SXin Lidef create_logger(level=logging.DEBUG, prefix=''): 116*9c5db199SXin Li """Creates a new logger. 117*9c5db199SXin Li 118*9c5db199SXin Li @param level: the desired logging level 119*9c5db199SXin Li @param prefix: the prefix to add to each log entry 120*9c5db199SXin Li """ 121*9c5db199SXin Li global _logging_service 122*9c5db199SXin Li if _logging_service is None: 123*9c5db199SXin Li _logging_service = _LoggingService() 124*9c5db199SXin Li atexit.register(_logging_service.shutdown) 125*9c5db199SXin Li return _logging_service.create_logger(level=level, prefix=prefix) 126*9c5db199SXin Li 127*9c5db199SXin Li 128*9c5db199SXin Lidef kill_or_log_returncode(*popens): 129*9c5db199SXin Li """Kills all the processes of the given Popens or logs the return code. 130*9c5db199SXin Li 131*9c5db199SXin Li @param popens: The Popens to be killed. 132*9c5db199SXin Li """ 133*9c5db199SXin Li for p in popens: 134*9c5db199SXin Li if p.poll() is None: 135*9c5db199SXin Li try: 136*9c5db199SXin Li p.kill() 137*9c5db199SXin Li except Exception as e: 138*9c5db199SXin Li logging.warning('failed to kill %d, %s', p.pid, e) 139*9c5db199SXin Li else: 140*9c5db199SXin Li logging.warning('command exit (pid=%d, rc=%d): %s', 141*9c5db199SXin Li p.pid, p.returncode, p.command) 142*9c5db199SXin Li 143*9c5db199SXin Li 144*9c5db199SXin Lidef wait_and_check_returncode(*popens): 145*9c5db199SXin Li """Wait for all the Popens and check the return code is 0. 146*9c5db199SXin Li 147*9c5db199SXin Li If the return code is not 0, it raises an RuntimeError. 148*9c5db199SXin Li 149*9c5db199SXin Li @param popens: The Popens to be checked. 150*9c5db199SXin Li """ 151*9c5db199SXin Li error_message = None 152*9c5db199SXin Li for p in popens: 153*9c5db199SXin Li if p.wait() != 0: 154*9c5db199SXin Li error_message = ('Command failed(%d, %d): %s' % 155*9c5db199SXin Li (p.pid, p.returncode, p.command)) 156*9c5db199SXin Li logging.error(error_message) 157*9c5db199SXin Li if error_message: 158*9c5db199SXin Li raise RuntimeError(error_message) 159*9c5db199SXin Li 160*9c5db199SXin Li 161*9c5db199SXin Lidef execute(args, stdin=None, stdout=TEE_TO_LOGS, stderr=TEE_TO_LOGS, 162*9c5db199SXin Li run_as=None): 163*9c5db199SXin Li """Executes a child command and wait for it. 164*9c5db199SXin Li 165*9c5db199SXin Li Returns the output from standard output if 'stdout' is subprocess.PIPE. 166*9c5db199SXin Li Raises RuntimeException if the return code of the child command is not 0. 167*9c5db199SXin Li 168*9c5db199SXin Li @param args: the command to be executed 169*9c5db199SXin Li @param stdin: the executed program's standard input 170*9c5db199SXin Li @param stdout: the executed program's standard output 171*9c5db199SXin Li @param stderr: the executed program's standard error 172*9c5db199SXin Li @param run_as: if not None, run the command as the given user 173*9c5db199SXin Li """ 174*9c5db199SXin Li ps = popen(args, stdin=stdin, stdout=stdout, stderr=stderr, 175*9c5db199SXin Li run_as=run_as) 176*9c5db199SXin Li out = ps.communicate()[0] if stdout == subprocess.PIPE else None 177*9c5db199SXin Li wait_and_check_returncode(ps) 178*9c5db199SXin Li return out 179*9c5db199SXin Li 180*9c5db199SXin Li 181*9c5db199SXin Lidef _run_as(user): 182*9c5db199SXin Li """Changes the uid and gid of the running process to be that of the 183*9c5db199SXin Li given user and configures its supplementary groups. 184*9c5db199SXin Li 185*9c5db199SXin Li Don't call this function directly, instead wrap it in a lambda and 186*9c5db199SXin Li pass that to the preexec_fn argument of subprocess.Popen. 187*9c5db199SXin Li 188*9c5db199SXin Li Example usage: 189*9c5db199SXin Li subprocess.Popen(..., preexec_fn=lambda: _run_as('chronos')) 190*9c5db199SXin Li 191*9c5db199SXin Li @param user: the user to run as 192*9c5db199SXin Li """ 193*9c5db199SXin Li pw = pwd.getpwnam(user) 194*9c5db199SXin Li os.setgid(pw.pw_gid) 195*9c5db199SXin Li os.initgroups(user, pw.pw_gid) 196*9c5db199SXin Li os.setuid(pw.pw_uid) 197*9c5db199SXin Li 198*9c5db199SXin Li 199*9c5db199SXin Lidef popen(args, stdin=None, stdout=TEE_TO_LOGS, stderr=TEE_TO_LOGS, env=None, 200*9c5db199SXin Li run_as=None): 201*9c5db199SXin Li """Returns a Popen object just as subprocess.Popen does but with the 202*9c5db199SXin Li executed command stored in Popen.command. 203*9c5db199SXin Li 204*9c5db199SXin Li @param args: the command to be executed 205*9c5db199SXin Li @param stdin: the executed program's standard input 206*9c5db199SXin Li @param stdout: the executed program's standard output 207*9c5db199SXin Li @param stderr: the executed program's standard error 208*9c5db199SXin Li @param env: the executed program's environment 209*9c5db199SXin Li @param run_as: if not None, run the command as the given user 210*9c5db199SXin Li """ 211*9c5db199SXin Li command_id = next(_command_serial_number) 212*9c5db199SXin Li prefix = '[%04d] ' % command_id 213*9c5db199SXin Li 214*9c5db199SXin Li if stdout is TEE_TO_LOGS: 215*9c5db199SXin Li stdout = create_logger(level=logging.DEBUG, prefix=prefix) 216*9c5db199SXin Li if stderr is TEE_TO_LOGS: 217*9c5db199SXin Li stderr = create_logger(level=logging.ERROR, prefix=prefix) 218*9c5db199SXin Li 219*9c5db199SXin Li command = ' '.join(pipes.quote(x) for x in args) 220*9c5db199SXin Li logging.info('%sRunning: %s', prefix, command) 221*9c5db199SXin Li 222*9c5db199SXin Li preexec_fn = None 223*9c5db199SXin Li if run_as is not None: 224*9c5db199SXin Li preexec_fn = lambda: _run_as(run_as) 225*9c5db199SXin Li 226*9c5db199SXin Li # The lock is required for http://crbug.com/323843. 227*9c5db199SXin Li with _popen_lock: 228*9c5db199SXin Li ps = subprocess.Popen(args, stdin=stdin, stdout=stdout, stderr=stderr, 229*9c5db199SXin Li env=env, preexec_fn=preexec_fn) 230*9c5db199SXin Li logging.info('%spid is %d', prefix, ps.pid) 231*9c5db199SXin Li ps.command_id = command_id 232*9c5db199SXin Li ps.command = command 233*9c5db199SXin Li return ps 234