You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
685 lines
24 KiB
685 lines
24 KiB
# Copyright 2011 The Chromium OS Authors. All rights reserved. |
|
# Use of this source code is governed by a BSD-style license that can be |
|
# found in the LICENSE file. |
|
"""Utilities to run commands in outside/inside chroot and on the board.""" |
|
|
|
from __future__ import print_function |
|
|
|
import getpass |
|
import os |
|
import re |
|
import select |
|
import signal |
|
import subprocess |
|
import sys |
|
import tempfile |
|
import time |
|
|
|
import logger |
|
import misc |
|
|
|
mock_default = False |
|
|
|
LOG_LEVEL = ('none', 'quiet', 'average', 'verbose') |
|
|
|
|
|
def InitCommandExecuter(mock=False): |
|
# pylint: disable=global-statement |
|
global mock_default |
|
# Whether to default to a mock command executer or not |
|
mock_default = mock |
|
|
|
|
|
def GetCommandExecuter(logger_to_set=None, mock=False, log_level='verbose'): |
|
# If the default is a mock executer, always return one. |
|
if mock_default or mock: |
|
return MockCommandExecuter(log_level, logger_to_set) |
|
else: |
|
return CommandExecuter(log_level, logger_to_set) |
|
|
|
|
|
class CommandExecuter(object): |
|
"""Provides several methods to execute commands on several environments.""" |
|
|
|
def __init__(self, log_level, logger_to_set=None): |
|
self.log_level = log_level |
|
if log_level == 'none': |
|
self.logger = None |
|
else: |
|
if logger_to_set is not None: |
|
self.logger = logger_to_set |
|
else: |
|
self.logger = logger.GetLogger() |
|
|
|
def GetLogLevel(self): |
|
return self.log_level |
|
|
|
def SetLogLevel(self, log_level): |
|
self.log_level = log_level |
|
|
|
def RunCommandGeneric(self, |
|
cmd, |
|
return_output=False, |
|
machine=None, |
|
username=None, |
|
command_terminator=None, |
|
command_timeout=None, |
|
terminated_timeout=10, |
|
print_to_console=True, |
|
except_handler=lambda p, e: None): |
|
"""Run a command. |
|
|
|
Returns triplet (returncode, stdout, stderr). |
|
""" |
|
|
|
cmd = str(cmd) |
|
|
|
if self.log_level == 'quiet': |
|
print_to_console = False |
|
|
|
if self.log_level == 'verbose': |
|
self.logger.LogCmd(cmd, machine, username, print_to_console) |
|
elif self.logger: |
|
self.logger.LogCmdToFileOnly(cmd, machine, username) |
|
if command_terminator and command_terminator.IsTerminated(): |
|
if self.logger: |
|
self.logger.LogError('Command was terminated!', print_to_console) |
|
return (1, '', '') |
|
|
|
if machine is not None: |
|
user = '' |
|
if username is not None: |
|
user = username + '@' |
|
cmd = "ssh -t -t %s%s -- '%s'" % (user, machine, cmd) |
|
|
|
# We use setsid so that the child will have a different session id |
|
# and we can easily kill the process group. This is also important |
|
# because the child will be disassociated from the parent terminal. |
|
# In this way the child cannot mess the parent's terminal. |
|
p = None |
|
try: |
|
p = subprocess.Popen(cmd, |
|
stdout=subprocess.PIPE, |
|
stderr=subprocess.PIPE, |
|
shell=True, |
|
preexec_fn=os.setsid) |
|
|
|
full_stdout = '' |
|
full_stderr = '' |
|
|
|
# Pull output from pipes, send it to file/stdout/string |
|
out = err = None |
|
pipes = [p.stdout, p.stderr] |
|
|
|
my_poll = select.poll() |
|
my_poll.register(p.stdout, select.POLLIN) |
|
my_poll.register(p.stderr, select.POLLIN) |
|
|
|
terminated_time = None |
|
started_time = time.time() |
|
|
|
while len(pipes): |
|
if command_terminator and command_terminator.IsTerminated(): |
|
os.killpg(os.getpgid(p.pid), signal.SIGTERM) |
|
if self.logger: |
|
self.logger.LogError('Command received termination request. ' |
|
'Killed child process group.', |
|
print_to_console) |
|
break |
|
|
|
l = my_poll.poll(100) |
|
for (fd, _) in l: |
|
if fd == p.stdout.fileno(): |
|
out = os.read(p.stdout.fileno(), 16384) |
|
if return_output: |
|
full_stdout += out |
|
if self.logger: |
|
self.logger.LogCommandOutput(out, print_to_console) |
|
if out == '': |
|
pipes.remove(p.stdout) |
|
my_poll.unregister(p.stdout) |
|
if fd == p.stderr.fileno(): |
|
err = os.read(p.stderr.fileno(), 16384) |
|
if return_output: |
|
full_stderr += err |
|
if self.logger: |
|
self.logger.LogCommandError(err, print_to_console) |
|
if err == '': |
|
pipes.remove(p.stderr) |
|
my_poll.unregister(p.stderr) |
|
|
|
if p.poll() is not None: |
|
if terminated_time is None: |
|
terminated_time = time.time() |
|
elif (terminated_timeout is not None and |
|
time.time() - terminated_time > terminated_timeout): |
|
if self.logger: |
|
self.logger.LogWarning('Timeout of %s seconds reached since ' |
|
'process termination.' % |
|
terminated_timeout, |
|
print_to_console) |
|
break |
|
|
|
if (command_timeout is not None and |
|
time.time() - started_time > command_timeout): |
|
os.killpg(os.getpgid(p.pid), signal.SIGTERM) |
|
if self.logger: |
|
self.logger.LogWarning('Timeout of %s seconds reached since process' |
|
'started. Killed child process group.' % |
|
command_timeout, print_to_console) |
|
break |
|
|
|
if out == err == '': |
|
break |
|
|
|
p.wait() |
|
if return_output: |
|
return (p.returncode, full_stdout, full_stderr) |
|
return (p.returncode, '', '') |
|
except BaseException as e: |
|
except_handler(p, e) |
|
raise |
|
|
|
def RunCommand(self, *args, **kwargs): |
|
"""Run a command. |
|
|
|
Takes the same arguments as RunCommandGeneric except for return_output. |
|
Returns a single value returncode. |
|
""" |
|
# Make sure that args does not overwrite 'return_output' |
|
assert len(args) <= 1 |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = False |
|
return self.RunCommandGeneric(*args, **kwargs)[0] |
|
|
|
def RunCommandWExceptionCleanup(self, *args, **kwargs): |
|
"""Run a command and kill process if exception is thrown. |
|
|
|
Takes the same arguments as RunCommandGeneric except for except_handler. |
|
Returns same as RunCommandGeneric. |
|
""" |
|
|
|
def KillProc(proc, _): |
|
if proc: |
|
os.killpg(os.getpgid(proc.pid), signal.SIGTERM) |
|
|
|
# Make sure that args does not overwrite 'except_handler' |
|
assert len(args) <= 8 |
|
assert 'except_handler' not in kwargs |
|
kwargs['except_handler'] = KillProc |
|
return self.RunCommandGeneric(*args, **kwargs) |
|
|
|
def RunCommandWOutput(self, *args, **kwargs): |
|
"""Run a command. |
|
|
|
Takes the same arguments as RunCommandGeneric except for return_output. |
|
Returns a triplet (returncode, stdout, stderr). |
|
""" |
|
# Make sure that args does not overwrite 'return_output' |
|
assert len(args) <= 1 |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = True |
|
return self.RunCommandGeneric(*args, **kwargs) |
|
|
|
def RemoteAccessInitCommand(self, chromeos_root, machine): |
|
command = '' |
|
command += '\nset -- --remote=' + machine |
|
command += '\n. ' + chromeos_root + '/src/scripts/common.sh' |
|
command += '\n. ' + chromeos_root + '/src/scripts/remote_access.sh' |
|
command += '\nTMP=$(mktemp -d)' |
|
command += "\nFLAGS \"$@\" || exit 1" |
|
command += '\nremote_access_init' |
|
return command |
|
|
|
def WriteToTempShFile(self, contents): |
|
handle, command_file = tempfile.mkstemp(prefix=os.uname()[1], suffix='.sh') |
|
os.write(handle, '#!/bin/bash\n') |
|
os.write(handle, contents) |
|
os.close(handle) |
|
return command_file |
|
|
|
def CrosLearnBoard(self, chromeos_root, machine): |
|
command = self.RemoteAccessInitCommand(chromeos_root, machine) |
|
command += '\nlearn_board' |
|
command += '\necho ${FLAGS_board}' |
|
retval, output, _ = self.RunCommandWOutput(command) |
|
if self.logger: |
|
self.logger.LogFatalIf(retval, 'learn_board command failed') |
|
elif retval: |
|
sys.exit(1) |
|
return output.split()[-1] |
|
|
|
def CrosRunCommandGeneric(self, |
|
cmd, |
|
return_output=False, |
|
machine=None, |
|
command_terminator=None, |
|
chromeos_root=None, |
|
command_timeout=None, |
|
terminated_timeout=10, |
|
print_to_console=True): |
|
"""Run a command on a ChromeOS box. |
|
|
|
Returns triplet (returncode, stdout, stderr). |
|
""" |
|
|
|
if self.log_level != 'verbose': |
|
print_to_console = False |
|
|
|
if self.logger: |
|
self.logger.LogCmd(cmd, print_to_console=print_to_console) |
|
self.logger.LogFatalIf(not machine, 'No machine provided!') |
|
self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!') |
|
else: |
|
if not chromeos_root or not machine: |
|
sys.exit(1) |
|
chromeos_root = os.path.expanduser(chromeos_root) |
|
|
|
# Write all commands to a file. |
|
command_file = self.WriteToTempShFile(cmd) |
|
retval = self.CopyFiles(command_file, |
|
command_file, |
|
dest_machine=machine, |
|
command_terminator=command_terminator, |
|
chromeos_root=chromeos_root, |
|
dest_cros=True, |
|
recursive=False, |
|
print_to_console=print_to_console) |
|
if retval: |
|
if self.logger: |
|
self.logger.LogError('Could not run remote command on machine.' |
|
' Is the machine up?') |
|
return (retval, '', '') |
|
|
|
command = self.RemoteAccessInitCommand(chromeos_root, machine) |
|
command += '\nremote_sh bash %s' % command_file |
|
command += "\nl_retval=$?; echo \"$REMOTE_OUT\"; exit $l_retval" |
|
retval = self.RunCommandGeneric(command, |
|
return_output, |
|
command_terminator=command_terminator, |
|
command_timeout=command_timeout, |
|
terminated_timeout=terminated_timeout, |
|
print_to_console=print_to_console) |
|
if return_output: |
|
connect_signature = ( |
|
'Initiating first contact with remote host\n' + 'Connection OK\n') |
|
connect_signature_re = re.compile(connect_signature) |
|
modded_retval = list(retval) |
|
modded_retval[1] = connect_signature_re.sub('', retval[1]) |
|
return modded_retval |
|
return retval |
|
|
|
def CrosRunCommand(self, *args, **kwargs): |
|
"""Run a command on a ChromeOS box. |
|
|
|
Takes the same arguments as CrosRunCommandGeneric except for return_output. |
|
Returns a single value returncode. |
|
""" |
|
# Make sure that args does not overwrite 'return_output' |
|
assert len(args) <= 1 |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = False |
|
return self.CrosRunCommandGeneric(*args, **kwargs)[0] |
|
|
|
def CrosRunCommandWOutput(self, *args, **kwargs): |
|
"""Run a command on a ChromeOS box. |
|
|
|
Takes the same arguments as CrosRunCommandGeneric except for return_output. |
|
Returns a triplet (returncode, stdout, stderr). |
|
""" |
|
# Make sure that args does not overwrite 'return_output' |
|
assert len(args) <= 1 |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = True |
|
return self.CrosRunCommandGeneric(*args, **kwargs) |
|
|
|
def ChrootRunCommandGeneric(self, |
|
chromeos_root, |
|
command, |
|
return_output=False, |
|
command_terminator=None, |
|
command_timeout=None, |
|
terminated_timeout=10, |
|
print_to_console=True, |
|
cros_sdk_options=''): |
|
"""Runs a command within the chroot. |
|
|
|
Returns triplet (returncode, stdout, stderr). |
|
""" |
|
|
|
if self.log_level != 'verbose': |
|
print_to_console = False |
|
|
|
if self.logger: |
|
self.logger.LogCmd(command, print_to_console=print_to_console) |
|
|
|
handle, command_file = tempfile.mkstemp( |
|
dir=os.path.join(chromeos_root, 'src/scripts'), |
|
suffix='.sh', |
|
prefix='in_chroot_cmd') |
|
os.write(handle, '#!/bin/bash\n') |
|
os.write(handle, command) |
|
os.write(handle, '\n') |
|
os.close(handle) |
|
|
|
os.chmod(command_file, 0777) |
|
|
|
# if return_output is set, run a dummy command first to make sure that |
|
# the chroot already exists. We want the final returned output to skip |
|
# the output from chroot creation steps. |
|
if return_output: |
|
ret = self.RunCommand('cd %s; cros_sdk %s -- true' % |
|
(chromeos_root, cros_sdk_options)) |
|
if ret: |
|
return (ret, '', '') |
|
|
|
# Run command_file inside the chroot, making sure that any "~" is expanded |
|
# by the shell inside the chroot, not outside. |
|
command = ("cd %s; cros_sdk %s -- bash -c '%s/%s'" % |
|
(chromeos_root, cros_sdk_options, misc.CHROMEOS_SCRIPTS_DIR, |
|
os.path.basename(command_file))) |
|
ret = self.RunCommandGeneric(command, |
|
return_output, |
|
command_terminator=command_terminator, |
|
command_timeout=command_timeout, |
|
terminated_timeout=terminated_timeout, |
|
print_to_console=print_to_console) |
|
os.remove(command_file) |
|
return ret |
|
|
|
def ChrootRunCommand(self, *args, **kwargs): |
|
"""Runs a command within the chroot. |
|
|
|
Takes the same arguments as ChrootRunCommandGeneric except for |
|
return_output. |
|
Returns a single value returncode. |
|
""" |
|
# Make sure that args does not overwrite 'return_output' |
|
assert len(args) <= 2 |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = False |
|
return self.ChrootRunCommandGeneric(*args, **kwargs)[0] |
|
|
|
def ChrootRunCommandWOutput(self, *args, **kwargs): |
|
"""Runs a command within the chroot. |
|
|
|
Takes the same arguments as ChrootRunCommandGeneric except for |
|
return_output. |
|
Returns a triplet (returncode, stdout, stderr). |
|
""" |
|
# Make sure that args does not overwrite 'return_output' |
|
assert len(args) <= 2 |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = True |
|
return self.ChrootRunCommandGeneric(*args, **kwargs) |
|
|
|
def RunCommands(self, |
|
cmdlist, |
|
machine=None, |
|
username=None, |
|
command_terminator=None): |
|
cmd = ' ;\n'.join(cmdlist) |
|
return self.RunCommand(cmd, |
|
machine=machine, |
|
username=username, |
|
command_terminator=command_terminator) |
|
|
|
def CopyFiles(self, |
|
src, |
|
dest, |
|
src_machine=None, |
|
dest_machine=None, |
|
src_user=None, |
|
dest_user=None, |
|
recursive=True, |
|
command_terminator=None, |
|
chromeos_root=None, |
|
src_cros=False, |
|
dest_cros=False, |
|
print_to_console=True): |
|
src = os.path.expanduser(src) |
|
dest = os.path.expanduser(dest) |
|
|
|
if recursive: |
|
src = src + '/' |
|
dest = dest + '/' |
|
|
|
if src_cros == True or dest_cros == True: |
|
if self.logger: |
|
self.logger.LogFatalIf(src_cros == dest_cros, |
|
'Only one of src_cros and desc_cros can ' |
|
'be True.') |
|
self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!') |
|
elif src_cros == dest_cros or not chromeos_root: |
|
sys.exit(1) |
|
if src_cros == True: |
|
cros_machine = src_machine |
|
else: |
|
cros_machine = dest_machine |
|
|
|
command = self.RemoteAccessInitCommand(chromeos_root, cros_machine) |
|
ssh_command = ( |
|
'ssh -p ${FLAGS_ssh_port}' + ' -o StrictHostKeyChecking=no' + |
|
' -o UserKnownHostsFile=$(mktemp)' + ' -i $TMP_PRIVATE_KEY') |
|
rsync_prefix = "\nrsync -r -e \"%s\" " % ssh_command |
|
if dest_cros == True: |
|
command += rsync_prefix + '%s root@%s:%s' % (src, dest_machine, dest) |
|
return self.RunCommand(command, |
|
machine=src_machine, |
|
username=src_user, |
|
command_terminator=command_terminator, |
|
print_to_console=print_to_console) |
|
else: |
|
command += rsync_prefix + 'root@%s:%s %s' % (src_machine, src, dest) |
|
return self.RunCommand(command, |
|
machine=dest_machine, |
|
username=dest_user, |
|
command_terminator=command_terminator, |
|
print_to_console=print_to_console) |
|
|
|
if dest_machine == src_machine: |
|
command = 'rsync -a %s %s' % (src, dest) |
|
else: |
|
if src_machine is None: |
|
src_machine = os.uname()[1] |
|
src_user = getpass.getuser() |
|
command = 'rsync -a %s@%s:%s %s' % (src_user, src_machine, src, dest) |
|
return self.RunCommand(command, |
|
machine=dest_machine, |
|
username=dest_user, |
|
command_terminator=command_terminator, |
|
print_to_console=print_to_console) |
|
|
|
def RunCommand2(self, |
|
cmd, |
|
cwd=None, |
|
line_consumer=None, |
|
timeout=None, |
|
shell=True, |
|
join_stderr=True, |
|
env=None, |
|
except_handler=lambda p, e: None): |
|
"""Run the command with an extra feature line_consumer. |
|
|
|
This version allow developers to provide a line_consumer which will be |
|
fed execution output lines. |
|
|
|
A line_consumer is a callback, which is given a chance to run for each |
|
line the execution outputs (either to stdout or stderr). The |
|
line_consumer must accept one and exactly one dict argument, the dict |
|
argument has these items - |
|
'line' - The line output by the binary. Notice, this string includes |
|
the trailing '\n'. |
|
'output' - Whether this is a stdout or stderr output, values are either |
|
'stdout' or 'stderr'. When join_stderr is True, this value |
|
will always be 'output'. |
|
'pobject' - The object used to control execution, for example, call |
|
pobject.kill(). |
|
|
|
Note: As this is written, the stdin for the process executed is |
|
not associated with the stdin of the caller of this routine. |
|
|
|
Args: |
|
cmd: Command in a single string. |
|
cwd: Working directory for execution. |
|
line_consumer: A function that will ba called by this function. See above |
|
for details. |
|
timeout: terminate command after this timeout. |
|
shell: Whether to use a shell for execution. |
|
join_stderr: Whether join stderr to stdout stream. |
|
env: Execution environment. |
|
except_handler: Callback for when exception is thrown during command |
|
execution. Passed process object and exception. |
|
|
|
Returns: |
|
Execution return code. |
|
|
|
Raises: |
|
child_exception: if fails to start the command process (missing |
|
permission, no such file, etc) |
|
""" |
|
|
|
class StreamHandler(object): |
|
"""Internal utility class.""" |
|
|
|
def __init__(self, pobject, fd, name, line_consumer): |
|
self._pobject = pobject |
|
self._fd = fd |
|
self._name = name |
|
self._buf = '' |
|
self._line_consumer = line_consumer |
|
|
|
def read_and_notify_line(self): |
|
t = os.read(fd, 1024) |
|
self._buf = self._buf + t |
|
self.notify_line() |
|
|
|
def notify_line(self): |
|
p = self._buf.find('\n') |
|
while p >= 0: |
|
self._line_consumer(line=self._buf[:p + 1], |
|
output=self._name, |
|
pobject=self._pobject) |
|
if p < len(self._buf) - 1: |
|
self._buf = self._buf[p + 1:] |
|
p = self._buf.find('\n') |
|
else: |
|
self._buf = '' |
|
p = -1 |
|
break |
|
|
|
def notify_eos(self): |
|
# Notify end of stream. The last line may not end with a '\n'. |
|
if self._buf != '': |
|
self._line_consumer(line=self._buf, |
|
output=self._name, |
|
pobject=self._pobject) |
|
self._buf = '' |
|
|
|
if self.log_level == 'verbose': |
|
self.logger.LogCmd(cmd) |
|
elif self.logger: |
|
self.logger.LogCmdToFileOnly(cmd) |
|
|
|
# We use setsid so that the child will have a different session id |
|
# and we can easily kill the process group. This is also important |
|
# because the child will be disassociated from the parent terminal. |
|
# In this way the child cannot mess the parent's terminal. |
|
pobject = None |
|
try: |
|
pobject = subprocess.Popen( |
|
cmd, |
|
cwd=cwd, |
|
bufsize=1024, |
|
env=env, |
|
shell=shell, |
|
universal_newlines=True, |
|
stdout=subprocess.PIPE, |
|
stderr=subprocess.STDOUT if join_stderr else subprocess.PIPE, |
|
preexec_fn=os.setsid) |
|
|
|
# We provide a default line_consumer |
|
if line_consumer is None: |
|
line_consumer = lambda **d: None |
|
start_time = time.time() |
|
poll = select.poll() |
|
outfd = pobject.stdout.fileno() |
|
poll.register(outfd, select.POLLIN | select.POLLPRI) |
|
handlermap = {outfd: |
|
StreamHandler(pobject, outfd, 'stdout', line_consumer)} |
|
if not join_stderr: |
|
errfd = pobject.stderr.fileno() |
|
poll.register(errfd, |
|
select.POLLIN | select.POLLPRI) |
|
handlermap[errfd] = StreamHandler(pobject, |
|
errfd, |
|
'stderr', |
|
line_consumer) |
|
while len(handlermap): |
|
readables = poll.poll(300) |
|
for (fd, evt) in readables: |
|
handler = handlermap[fd] |
|
if evt & (select.POLLPRI | select.POLLIN): |
|
handler.read_and_notify_line() |
|
elif evt & (select.POLLHUP | select.POLLERR | select.POLLNVAL): |
|
handler.notify_eos() |
|
poll.unregister(fd) |
|
del handlermap[fd] |
|
|
|
if timeout is not None and (time.time() - start_time > timeout): |
|
os.killpg(os.getpgid(pobject.pid), signal.SIGTERM) |
|
|
|
return pobject.wait() |
|
except BaseException as e: |
|
except_handler(pobject, e) |
|
raise |
|
|
|
|
|
class MockCommandExecuter(CommandExecuter): |
|
"""Mock class for class CommandExecuter.""" |
|
|
|
def __init__(self, log_level, logger_to_set=None): |
|
super(MockCommandExecuter, self).__init__(log_level, logger_to_set) |
|
|
|
def RunCommandGeneric(self, |
|
cmd, |
|
return_output=False, |
|
machine=None, |
|
username=None, |
|
command_terminator=None, |
|
command_timeout=None, |
|
terminated_timeout=10, |
|
print_to_console=True, |
|
except_handler=lambda p, e: None): |
|
assert not command_timeout |
|
cmd = str(cmd) |
|
if machine is None: |
|
machine = 'localhost' |
|
if username is None: |
|
username = 'current' |
|
logger.GetLogger().LogCmd('(Mock) ' + cmd, machine, username, |
|
print_to_console) |
|
return (0, '', '') |
|
|
|
def RunCommand(self, *args, **kwargs): |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = False |
|
return self.RunCommandGeneric(*args, **kwargs)[0] |
|
|
|
def RunCommandWOutput(self, *args, **kwargs): |
|
assert 'return_output' not in kwargs |
|
kwargs['return_output'] = True |
|
return self.RunCommandGeneric(*args, **kwargs) |
|
|
|
|
|
class CommandTerminator(object): |
|
"""Object to request termination of a command in execution.""" |
|
|
|
def __init__(self): |
|
self.terminated = False |
|
|
|
def Terminate(self): |
|
self.terminated = True |
|
|
|
def IsTerminated(self): |
|
return self.terminated
|
|
|