1# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import contextlib
6import getpass
7import subprocess
8import os
9
10import common
11from autotest_lib.server.hosts import ssh_host
12from autotest_lib.client.common_lib import error
13from autotest_lib.client.common_lib import global_config
14from autotest_lib.client.common_lib import utils
15from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
16
17
18@contextlib.contextmanager
19def chdir(dirname=None):
20    """A context manager to help change directories.
21
22    Will chdir into the provided dirname for the lifetime of the context and
23    return to cwd thereafter.
24
25    @param dirname: The dirname to chdir into.
26    """
27    curdir = os.getcwd()
28    try:
29        if dirname is not None:
30            os.chdir(dirname)
31        yield
32    finally:
33        os.chdir(curdir)
34
35
36def local_runner(cmd, stream_output=False):
37    """
38    Runs a command on the local system as the current user.
39
40    @param cmd: The command to run.
41    @param stream_output: If True, streams the stdout of the process.
42
43    @returns: The output of cmd.
44    @raises CalledProcessError: If there was a non-0 return code.
45    """
46    if not stream_output:
47        return subprocess.check_output(cmd, shell=True)
48    proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
49    while proc.poll() is None:
50        print proc.stdout.readline().rstrip('\n')
51
52
53_host_objects = {}
54
55def host_object_runner(host, **kwargs):
56    """
57    Returns a function that returns the output of running a command via a host
58    object.
59
60    @param host: The host to run a command on.
61    @returns: A function that can invoke a command remotely.
62    """
63    try:
64        host_object = _host_objects[host]
65    except KeyError:
66        username = global_config.global_config.get_config_value(
67                'CROS', 'infrastructure_user')
68        host_object = ssh_host.SSHHost(host, user=username)
69        _host_objects[host] = host_object
70
71    def runner(cmd):
72        """
73        Runs a command via a host object on the enclosed host.  Translates
74        host.run errors to the subprocess equivalent to expose a common API.
75
76        @param cmd: The command to run.
77        @returns: The output of cmd.
78        @raises CalledProcessError: If there was a non-0 return code.
79        """
80        try:
81            return host_object.run(cmd).stdout
82        except error.AutotestHostRunError as e:
83            exit_status = e.result_obj.exit_status
84            command = e.result_obj.command
85            raise subprocess.CalledProcessError(exit_status, command)
86    return runner
87
88
89def googlesh_runner(host, **kwargs):
90    """
91    Returns a function that return the output of running a command via shelling
92    out to `googlesh`.
93
94    @param host: The host to run a command on
95    @returns: A function that can invoke a command remotely.
96    """
97    def runner(cmd):
98        """
99        Runs a command via googlesh on the enclosed host.
100
101        @param cmd: The command to run.
102        @returns: The output of cmd.
103        @raises CalledProcessError: If there was a non-0 return code.
104        """
105        out = subprocess.check_output(['googlesh', '-s', '-uchromeos-test',
106                                       '-m%s' % host, '%s' % cmd])
107        return out
108    return runner
109
110
111def execute_command(host, cmd, **kwargs):
112    """
113    Executes a command on the host `host`.  This an optimization that if
114    we're already chromeos-test, we can just ssh to the machine in question.
115    Or if we're local, we don't have to ssh at all.
116
117    @param host: The hostname to execute the command on.
118    @param cmd: The command to run.  Special shell syntax (such as pipes)
119                is allowed.
120    @param kwargs: Key word arguments for the runner functions.
121    @returns: The output of the command.
122    """
123    if utils.is_localhost(host):
124        runner = local_runner
125    elif getpass.getuser() == 'chromeos-test':
126        runner = host_object_runner(host)
127    else:
128        runner = googlesh_runner(host)
129
130    return runner(cmd, **kwargs)
131
132
133def _csv_to_list(s):
134    """
135    Converts a list seperated by commas into a list of strings.
136
137    >>> _csv_to_list('')
138    []
139    >>> _csv_to_list('one')
140    ['one']
141    >>> _csv_to_list('one, two,three')
142    ['one', 'two', 'three']
143    """
144    return [x.strip() for x in s.split(',') if x]
145
146
147# The goal with these functions is to give you a list of hosts that are valid
148# arguments to ssh.  Note that this only really works since our instances use
149# names that are findable by our default /etc/resolv.conf `search` domains,
150# because all of our instances have names under .corp
151def sam_servers():
152    """
153    Generate a list of all scheduler/afe instances of autotest.
154
155    Note that we don't include the mysql database host if the database is split
156    from the rest of the system.
157    """
158    sams_config = global_config.global_config.get_config_value(
159            'CROS', 'sam_instances', default='')
160    sams = _csv_to_list(sams_config)
161    return set(sams)
162
163
164def extra_servers():
165    """
166    Servers that have an autotest checkout in /usr/local/autotest, but aren't
167    in any other list.
168
169    @returns: A set of hosts.
170    """
171    servers = global_config.global_config.get_config_value(
172                'CROS', 'extra_servers', default='')
173    return set(_csv_to_list(servers))
174
175
176def test_instance():
177    """
178    A server that is set up to run tests of the autotest infrastructure.
179
180    @returns: A hostname
181    """
182    server = global_config.global_config.get_config_value(
183                'CROS', 'test_instance', default='')
184    return server
185
186
187# The most reliable way to pull information about the state of the lab is to
188# look at the global/shadow config on each server.  The best way to do this is
189# via the global_config module.  Therefore, we invoke python on the remote end
190# to call global_config to get whatever values we want.
191_VALUE_FROM_CONFIG = '''
192cd /usr/local/autotest
193python -c "
194import common
195from autotest_lib.client.common_lib import global_config
196print global_config.global_config.get_config_value(
197  '%s', '%s', default='')
198"
199'''
200# There's possibly cheaper ways to do some of this, for example, we could scrape
201# instance:13467 for the list of drones, but this way you can get the list of
202# drones that is what should/will be running, and not what the scheduler thinks
203# is running.  (It could have kicked one out, or we could be bringing a new one
204# into rotation.)  So scraping the config on remote servers, while slow, gives
205# us consistent logical results.
206
207
208def _scrape_from_instances(section, key):
209    sams = sam_servers()
210    all_servers = set()
211    for sam in sams:
212        servers_csv = execute_command(sam, _VALUE_FROM_CONFIG % (section, key))
213        servers = _csv_to_list(servers_csv)
214        for server in servers:
215            if server == 'localhost':
216                all_servers.add(sam)
217            else:
218                all_servers.add(server)
219    return all_servers
220
221
222def database_servers():
223    """
224    Generate a list of all database servers running for instances of autotest.
225
226    @returns: An iterable of all hosts.
227    """
228    return _scrape_from_instances('AUTOTEST_WEB', 'host')
229
230
231def drone_servers():
232    """
233    Generate a list of all drones used by all instances of autotest in
234    production.
235
236    @returns: An iterable of drone servers.
237    """
238    return _scrape_from_instances('SCHEDULER', 'drones')
239
240
241def devserver_servers():
242    """
243    Generate a list of all devservers.
244
245    @returns: An iterable of all hosts.
246    """
247    zone = global_config.global_config.get_config_value(
248            'CLIENT', 'dns_zone')
249    servers = _scrape_from_instances('CROS', 'dev_server_hosts')
250    # The default text we get back here isn't something you can ssh into unless
251    # you've set up your /etc/resolve.conf to automatically try .cros, so we
252    # append the zone to try and make this more in line with everything else.
253    return set([server+'.'+zone for server in servers])
254
255
256def shard_servers():
257    """
258    Generate a list of all shard servers.
259
260    @returns: An iterable of all shard servers.
261    """
262    shard_hostnames = set()
263    sams = sam_servers()
264    for sam in sams:
265        afe = frontend_wrappers.RetryingAFE(server=sam)
266        shards = afe.run('get_shards')
267        for shard in shards:
268            shard_hostnames.add(shard['hostname'])
269
270    return list(shard_hostnames)
271