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, will be stdout and stderr.
44    @raises CalledProcessError: If there was a non-0 return code.
45    """
46    print 'Running command: %s' % cmd
47    proc = subprocess.Popen(
48        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
49    if stream_output:
50        output = ''
51        for newline in iter(proc.stdout.readline, ''):
52            output += newline
53            print newline.rstrip(os.linesep)
54    else:
55        output = proc.communicate()[0]
56
57    return_code = proc.wait()
58    if return_code !=0:
59        print "ERROR: '%s' failed with error:\n%s" % (cmd, output)
60        raise subprocess.CalledProcessError(return_code, cmd, output[:1024])
61    return output
62
63
64_host_objects = {}
65
66def host_object_runner(host, **kwargs):
67    """
68    Returns a function that returns the output of running a command via a host
69    object.
70
71    @param host: The host to run a command on.
72    @returns: A function that can invoke a command remotely.
73    """
74    try:
75        host_object = _host_objects[host]
76    except KeyError:
77        username = global_config.global_config.get_config_value(
78                'CROS', 'infrastructure_user')
79        host_object = ssh_host.SSHHost(host, user=username)
80        _host_objects[host] = host_object
81
82    def runner(cmd):
83        """
84        Runs a command via a host object on the enclosed host.  Translates
85        host.run errors to the subprocess equivalent to expose a common API.
86
87        @param cmd: The command to run.
88        @returns: The output of cmd.
89        @raises CalledProcessError: If there was a non-0 return code.
90        """
91        try:
92            return host_object.run(cmd).stdout
93        except error.AutotestHostRunError as e:
94            exit_status = e.result_obj.exit_status
95            command = e.result_obj.command
96            raise subprocess.CalledProcessError(exit_status, command)
97    return runner
98
99
100def googlesh_runner(host, **kwargs):
101    """
102    Returns a function that return the output of running a command via shelling
103    out to `googlesh`.
104
105    @param host: The host to run a command on
106    @returns: A function that can invoke a command remotely.
107    """
108    def runner(cmd):
109        """
110        Runs a command via googlesh on the enclosed host.
111
112        @param cmd: The command to run.
113        @returns: The output of cmd.
114        @raises CalledProcessError: If there was a non-0 return code.
115        """
116        out = subprocess.check_output(['googlesh', '-s', '-uchromeos-test',
117                                       '-m%s' % host, '%s' % cmd],
118                                      stderr=subprocess.STDOUT)
119        return out
120    return runner
121
122
123def execute_command(host, cmd, **kwargs):
124    """
125    Executes a command on the host `host`.  This an optimization that if
126    we're already chromeos-test, we can just ssh to the machine in question.
127    Or if we're local, we don't have to ssh at all.
128
129    @param host: The hostname to execute the command on.
130    @param cmd: The command to run.  Special shell syntax (such as pipes)
131                is allowed.
132    @param kwargs: Key word arguments for the runner functions.
133    @returns: The output of the command.
134    """
135    if utils.is_localhost(host):
136        runner = local_runner
137    elif getpass.getuser() == 'chromeos-test':
138        runner = host_object_runner(host)
139    else:
140        runner = googlesh_runner(host)
141
142    return runner(cmd, **kwargs)
143