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.

717 lines
24 KiB

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 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
from cros_utils import logger
mock_default = False
CHROMEOS_SCRIPTS_DIR = '/mnt/host/source/src/scripts'
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,
env=None,
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:
# pylint: disable=bad-option-value, subprocess-popen-preexec-fn
p = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
shell=True,
preexec_fn=os.setsid,
executable='/bin/bash',
env=env)
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 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).decode('utf8')
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).decode('utf8')
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 err:
except_handler(p, err)
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, port=None):
command = ''
command += '\nset -- --remote=' + machine
if port:
command += ' --ssh_port=' + port
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):
with tempfile.NamedTemporaryFile(
'w', encoding='utf-8', delete=False, prefix=os.uname()[1],
suffix='.sh') as f:
f.write('#!/bin/bash\n')
f.write(contents)
f.flush()
return f.name
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)
port = None
if ':' in machine:
machine, port = machine.split(':')
# Write all commands to a file.
command_file = self.WriteToTempShFile(cmd)
retval = self.CopyFiles(
command_file,
command_file,
dest_machine=machine,
dest_port=port,
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, port)
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='',
env=None):
"""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)
with tempfile.NamedTemporaryFile(
'w',
encoding='utf-8',
delete=False,
dir=os.path.join(chromeos_root, 'src/scripts'),
suffix='.sh',
prefix='in_chroot_cmd') as f:
f.write('#!/bin/bash\n')
f.write(command)
f.write('\n')
f.flush()
command_file = f.name
os.chmod(command_file, 0o777)
# if return_output is set, run a test 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),
env=env)
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, 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,
env=env)
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,
src_port=None,
dest_machine=None,
dest_port=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 or dest_cros:
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:
cros_machine = src_machine
cros_port = src_port
host_machine = dest_machine
host_user = dest_user
else:
cros_machine = dest_machine
cros_port = dest_port
host_machine = src_machine
host_user = src_user
command = self.RemoteAccessInitCommand(chromeos_root, cros_machine,
cros_port)
ssh_command = ('ssh -o StrictHostKeyChecking=no' +
' -o UserKnownHostsFile=$(mktemp)' +
' -i $TMP_PRIVATE_KEY')
if cros_port:
ssh_command += ' -p %s' % cros_port
rsync_prefix = '\nrsync -r -e "%s" ' % ssh_command
if dest_cros:
command += rsync_prefix + '%s root@%s:%s' % (src, cros_machine, dest)
else:
command += rsync_prefix + 'root@%s:%s %s' % (cros_machine, src, dest)
return self.RunCommand(
command,
machine=host_machine,
username=host_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:
# pylint: disable=bad-option-value, subprocess-popen-preexec-fn
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 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 err:
except_handler(pobject, err)
raise
class MockCommandExecuter(CommandExecuter):
"""Mock class for class CommandExecuter."""
def RunCommandGeneric(self,
cmd,
return_output=False,
machine=None,
username=None,
command_terminator=None,
command_timeout=None,
terminated_timeout=10,
print_to_console=True,
env=None,
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