site_utils/lib/infra.py

# Copyright (c) 2014 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.

import contextlib
import getpass
import subprocess
import os

import common
from autotest_lib.server.hosts import ssh_host
from autotest_lib.client.common_lib import error
from autotest_lib.client.common_lib import global_config
from autotest_lib.client.common_lib import utils
from autotest_lib.server.cros.dynamic_suite import frontend_wrappers


@contextlib.contextmanager
def chdir(dirname=None):
    """A context manager to help change directories.

    Will chdir into the provided dirname for the lifetime of the context and
    return to cwd thereafter.

    @param dirname: The dirname to chdir into.
    """
    curdir = os.getcwd()
    try:
        if dirname is not None:
            os.chdir(dirname)
        yield
    finally:
        os.chdir(curdir)


def local_runner(cmd, stream_output=False):
    """
    Runs a command on the local system as the current user.

    @param cmd: The command to run.
    @param stream_output: If True, streams the stdout of the process.

    @returns: The output of cmd, will be stdout and stderr.
    @raises CalledProcessError: If there was a non-0 return code.
    """
    print 'Running command: %s' % cmd
    proc = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if stream_output:
        output = ''
        for newline in iter(proc.stdout.readline, ''):
            output += newline
            print newline.rstrip(os.linesep)
    else:
        output = proc.communicate()[0]

    return_code = proc.wait()
    if return_code !=0:
        print "ERROR: '%s' failed with error:\n%s" % (cmd, output)
        raise subprocess.CalledProcessError(return_code, cmd, output[:1024])
    return output


_host_objects = {}

def host_object_runner(host, **kwargs):
    """
    Returns a function that returns the output of running a command via a host
    object.

    @param host: The host to run a command on.
    @returns: A function that can invoke a command remotely.
    """
    try:
        host_object = _host_objects[host]
    except KeyError:
        username = global_config.global_config.get_config_value(
                'CROS', 'infrastructure_user')
        host_object = ssh_host.SSHHost(host, user=username)
        _host_objects[host] = host_object

    def runner(cmd):
        """
        Runs a command via a host object on the enclosed host.  Translates
        host.run errors to the subprocess equivalent to expose a common API.

        @param cmd: The command to run.
        @returns: The output of cmd.
        @raises CalledProcessError: If there was a non-0 return code.
        """
        try:
            return host_object.run(cmd).stdout
        except error.AutotestHostRunError as e:
            exit_status = e.result_obj.exit_status
            command = e.result_obj.command
            raise subprocess.CalledProcessError(exit_status, command)
    return runner


def googlesh_runner(host, **kwargs):
    """
    Returns a function that return the output of running a command via shelling
    out to `googlesh`.

    @param host: The host to run a command on
    @returns: A function that can invoke a command remotely.
    """
    def runner(cmd):
        """
        Runs a command via googlesh on the enclosed host.

        @param cmd: The command to run.
        @returns: The output of cmd.
        @raises CalledProcessError: If there was a non-0 return code.
        """
        out = subprocess.check_output(['googlesh', '-s', '-uchromeos-test',
                                       '-m%s' % host, '%s' % cmd],
                                      stderr=subprocess.STDOUT)
        return out
    return runner


def execute_command(host, cmd, **kwargs):
    """
    Executes a command on the host `host`.  This an optimization that if
    we're already chromeos-test, we can just ssh to the machine in question.
    Or if we're local, we don't have to ssh at all.

    @param host: The hostname to execute the command on.
    @param cmd: The command to run.  Special shell syntax (such as pipes)
                is allowed.
    @param kwargs: Key word arguments for the runner functions.
    @returns: The output of the command.
    """
    if utils.is_localhost(host):
        runner = local_runner
    elif getpass.getuser() == 'chromeos-test':
        runner = host_object_runner(host)
    else:
        runner = googlesh_runner(host)

    return runner(cmd, **kwargs)