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 getpass
import socket
import subprocess

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


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

    @param cmd: The command to run.
    @returns: The output of cmd.
    @raises CalledProcessError: If there was a non-0 return code.
    """
    return subprocess.check_output(cmd, shell=True)


_host_objects = {}

def host_object_runner(host):
    """
    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 become_runner(host):
    """
    Returns a function that return the output of running a command via shelling
    out to `become`.

    @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 become 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(['become', 'chromeos-test@%s' % host,
                                       '--', cmd])
        return out
    return runner


def execute_command(host, cmd):
    """
    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.
    @returns: The output of the command.
    """
    if (socket.gethostbyaddr(host)[0] == socket.gethostname()
        or host == 'localhost'):
        runner = local_runner
    elif getpass.getuser() == 'chromeos-test':
        runner = host_object_runner(host)
    else:
        runner = become_runner(host)

    return runner(cmd)


def _csv_to_list(s):
    """
    Converts a list seperated by commas into a list of strings.

    >>> _csv_to_list('')
    []
    >>> _csv_to_list('one')
    ['one']
    >>> _csv_to_list('one, two,three')
    ['one', 'two', 'three']
    """
    return [x.strip() for x in s.split(',') if x]


# The goal with these functions is to give you a list of hosts that are valid
# arguments to ssh.  Note that this only really works since our instances use
# names that are findable by our default /etc/resolv.conf `search` domains,
# because all of our instances have names under .corp
def sam_servers():
    """
    Generate a list of all scheduler/afe instances of autotest.

    Note that we don't include the mysql database host if the database is split
    from the rest of the system.
    """
    sams_config = global_config.global_config.get_config_value(
            'SERVER', 'sam_instances', default='')
    sams = _csv_to_list(sams_config)
    return sams


# The most reliable way to pull information about the state of the lab is to
# look at the global/shadow config on each server.  The best way to do this is
# via the global_config module.  Therefore, we invoke python on the remote end
# to call global_config to get whatever values we want.
_VALUE_FROM_CONFIG = '''
cd /usr/local/autotest
python -c "
import common
from autotest_lib.client.common_lib import global_config
print global_config.global_config.get_config_value(
  '%s', '%s', default='')
"
'''
# There's possibly cheaper ways to do some of this, for example, we could scrape
# instance:13467 for the list of drones, but this way you can get the list of
# drones that is what should/will be running, and not what the scheduler thinks
# is running.  (It could have kicked one out, or we could be bringing a new one
# into rotation.)  So scraping the config on remote servers, while slow, gives
# us consistent logical results.


def _scrape_from_instances(section, key):
    sams = sam_servers()
    all_servers = set()
    for sam in sams:
        servers_csv = execute_command(sam, _VALUE_FROM_CONFIG % (section, key))
        servers = _csv_to_list(servers_csv)
        for server in servers:
            if server == 'localhost':
                all_servers.add(sam)
            else:
                all_servers.add(server)
    return all_servers


def database_servers():
    """
    Generate a list of all database servers running for instances of autotest.

    @returns: An iterable of all hosts.
    """
    return _scrape_from_instances('AUTOTEST_WEB', 'host')


def drone_servers():
    """
    Generate a list of all drones used by all instances of autotest in
    production.

    @returns: An iterable of all hosts.
    """
    return _scrape_from_instances('SCHEDULER', 'drones')


def devserver_servers():
    """
    Generate a list of all devservers.

    @returns: An iterable of all hosts.
    """
    zone = global_config.global_config.get_config_value(
            'CROS', 'dns_zone')
    servers = _scrape_from_instances('CROS', 'dev_server_hosts')
    # The default text we get back here isn't something you can ssh into unless
    # you've set up your /etc/resolve.conf to automatically try .cros, so we
    # append the zone to try and make this more in line with everything else.
    return set([server+'.'+zone for server in servers])