diff options
Diffstat (limited to 'utils/command_executer.py')
-rw-r--r-- | utils/command_executer.py | 304 |
1 files changed, 304 insertions, 0 deletions
diff --git a/utils/command_executer.py b/utils/command_executer.py new file mode 100644 index 00000000..32d86c91 --- /dev/null +++ b/utils/command_executer.py @@ -0,0 +1,304 @@ +import getpass +import logger +import os +import pty +import re +import select +import subprocess +import sys +import tempfile +import time +import utils + +mock_default = False + + +def InitCommandExecuter(mock=False): + global mock_default + # Whether to default to a mock command executer or not + mock_default = mock + + +def GetCommandExecuter(logger_to_set=None, mock=False): + # If the default is a mock executer, always return one. + if mock_default or mock: + return MockCommandExecuter(logger_to_set) + else: + return CommandExecuter(logger_to_set) + + +class CommandExecuter: + def __init__(self, logger_to_set=None): + if logger_to_set is not None: + self.logger = logger_to_set + else: + self.logger = logger.GetLogger() + + def RunCommand(self, cmd, return_output=False, machine=None, + username=None, command_terminator=None, + command_timeout=None, + terminated_timeout=10): + """Run a command.""" + + cmd = str(cmd) + + self.logger.LogCmd(cmd, machine, username) + if command_terminator and command_terminator.IsTerminated(): + self.logger.LogError("Command was terminated!") + if return_output: + return [1, "", ""] + else: + 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) + + pty_fds = pty.openpty() + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + stdin=pty_fds[0], shell=True) + + full_stdout = "" + full_stderr = "" + + # Pull output from pipes, send it to file/stdout/string + out = err = None + pipes = [p.stdout, p.stderr] + + terminated_time = None + started_time = time.time() + + while len(pipes): + fds = select.select(pipes, [], [], 0.1) + if command_terminator and command_terminator.IsTerminated(): + self.RunCommand("sudo kill -9 " + str(p.pid)) + wait = p.wait() + self.logger.LogError("Command was terminated!") + if return_output: + return (p.wait, full_stdout, full_stderr) + else: + return wait + for fd in fds[0]: + if fd == p.stdout: + out = os.read(p.stdout.fileno(), 16384) + if return_output: + full_stdout += out + self.logger.LogCommandOutput(out) + if out == "": + pipes.remove(p.stdout) + if fd == p.stderr: + err = os.read(p.stderr.fileno(), 16384) + if return_output: + full_stderr += err + self.logger.LogCommandError(err) + if err == "": + pipes.remove(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): + m = ("Timeout of %s seconds reached since process termination." + % terminated_timeout) + self.logger.LogWarning(m) + break + + if (command_timeout is not None and + time.time() - started_time > command_timeout): + m = ("Timeout of %s seconds reached since process started." + % command_timeout) + self.logger.LogWarning(m) + self.RunCommand("kill %d || sudo kill %d || sudo kill -9 %d" % + (p.pid, p.pid, p.pid)) + break + + if out == err == "": + break + + p.wait() + os.close(pty_fds[0]) + os.close(pty_fds[1]) + if return_output: + return (p.returncode, full_stdout, full_stderr) + return p.returncode + + 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, err = self.RunCommand(command, True) + self.logger.LogFatalIf(retval, "learn_board command failed") + return output.split()[-1] + + def CrosRunCommand(self, cmd, return_output=False, machine=None, + username=None, command_terminator=None, chromeos_root=None, + command_timeout=None): + """Run a command on a chromeos box""" + self.logger.LogCmd(cmd) + self.logger.LogFatalIf(not machine, "No machine provided!") + self.logger.LogFatalIf(not chromeos_root, "chromeos_root not given!") + chromeos_root = os.path.expanduser(chromeos_root) + + # Write all commands to a file. + command_file = self.WriteToTempShFile(cmd) + self.CopyFiles(command_file, command_file, + dest_machine=machine, + command_terminator=command_terminator, + chromeos_root=chromeos_root, + dest_cros=True, + recursive=False) + + command = self.RemoteAccessInitCommand(chromeos_root, machine) + command += "\nremote_sh bash %s" % command_file + command += "\necho \"$REMOTE_OUT\"" + retval = self.RunCommand(command, return_output, + command_terminator=command_terminator, + command_timeout=command_timeout) + if return_output: + connect_signature = ("Initiating first contact with remote host\n" + + "Connection OK\n") + connect_signature_re = re.compile(connect_signature) + modded_return = [] + for r in retval: + modded_return.append(r) + modded_return[1] = connect_signature_re.sub("", modded_return[1]) + return modded_return + return retval + + def ChrootRunCommand(self, chromeos_root, command, return_output=False, + command_terminator=None): + self.logger.LogCmd(command) + + 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.close(handle) + + os.chmod(command_file, 0777) + + command = "cd %s; cros_sdk -- ./%s" % (chromeos_root, + os.path.basename(command_file)) + ret = self.RunCommand(command, return_output, + command_terminator=command_terminator) + os.remove(command_file) + return ret + + + def RunCommands(self, cmdlist, return_output=False, machine=None, + username=None, command_terminator=None): + cmd = " ;\n" .join(cmdlist) + return self.RunCommand(cmd, return_output, machine, username, + 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): + src = os.path.expanduser(src) + dest = os.path.expanduser(dest) + + if recursive: + src = src + "/" + dest = dest + "/" + + if src_cros == True or dest_cros == True: + self.logger.LogFatalIf(not (src_cros ^ dest_cros), "Only one of src_cros " + "and desc_cros can be non-null.") + self.logger.LogFatalIf(not chromeos_root, "chromeos_root not given!") + if src_cros == True: + cros_machine = src_machine + else: + cros_machine = dest_machine + + command = self.RemoteAccessInitCommand(chromeos_root, cros_machine) + src_parent, src_child = utils.GetRoot(src) + dest_parent, dest_child = utils.GetRoot(dest) + 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) + 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) + + + 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) + + +class MockCommandExecuter(CommandExecuter): + def __init__(self, logger_to_set=None): + if logger is not None: + self.logger = logger_to_set + else: + self.logger = logger.GetLogger() + + def RunCommand(self, cmd, return_output=False, machine=None, username=None, + command_terminator=None): + cmd = str(cmd) + if machine is None: + machine = "localhost" + if username is None: + username = "current" + logger.GetLogger().LogCmd("(Mock) " + cmd, machine, username) + return 0 + + +class CommandTerminator: + def __init__(self): + self.terminated = False + + def Terminate(self): + self.terminated = True + + def IsTerminated(self): + return self.terminated |