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 5 import contextlib 6 import getpass 7 import subprocess 8 import os 9 10 import common 11 from autotest_lib.server.hosts import ssh_host 12 from autotest_lib.client.common_lib import error 13 from autotest_lib.client.common_lib import global_config 14 from autotest_lib.client.common_lib import utils 15 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 16 17 18 @contextlib.contextmanager 19 def 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 36 def 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 66 def 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 100 def 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 123 def 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