Home | History | Annotate | Download | only in lib 
      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