1 # Copyright 2015 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 logging 6 import os 7 import string 8 9 import common 10 from chromite.lib import gce 11 12 from autotest_lib.client.common_lib import error 13 from autotest_lib.client.common_lib import lsbrelease_utils 14 from autotest_lib.client.cros import constants as client_constants 15 from autotest_lib.server.hosts import abstract_ssh 16 17 SSH_KEYS_METADATA_KEY = "sshKeys" 18 TMP_DIR='/usr/local/tmp' 19 20 def extract_arguments(args_dict): 21 """Extract GCE-specific arguments from arguments dictionary. 22 23 @param args_dict: dictionary of all arguments supplied to the test. 24 """ 25 26 return {k: v for k, v in args_dict.items() 27 if k in ('gce_project', 'gce_instance', 28 'gce_zone', 'gce_key_file')} 29 30 31 class GceHost(abstract_ssh.AbstractSSHHost): 32 """GCE-specific subclass of Host.""" 33 34 def _initialize(self, hostname, gce_args=None, 35 *args, **dargs): 36 """Initializes this instance of GceHost. 37 38 @param hostname: the hostnname to be passed down to AbstractSSHHost. 39 @param gce_args: GCE-specific arguments extracted using 40 extract_arguments(). 41 """ 42 super(GceHost, self)._initialize(hostname=hostname, 43 *args, **dargs) 44 45 if gce_args: 46 self._gce_project = gce_args['gce_project'] 47 self._gce_zone = gce_args['gce_zone'] 48 self._gce_instance = gce_args['gce_instance'] 49 self._gce_key_file = gce_args['gce_key_file'] 50 else: 51 logging.warning("No GCE flags provided, calls to GCE API will fail") 52 return 53 54 self.gce = gce.GceContext.ForServiceAccountThreadSafe( 55 self._gce_project, self._gce_zone, self._gce_key_file) 56 57 @staticmethod 58 def check_host(host, timeout=10): 59 """ 60 Check if the given host is running on GCE. 61 62 @param host: An ssh host representing a device. 63 @param timeout: The timeout for the run command. 64 65 @return: True if the host is running on GCE. 66 """ 67 try: 68 result = host.run( 69 'grep CHROMEOS_RELEASE_BOARD /etc/lsb-release', 70 timeout=timeout) 71 return lsbrelease_utils.is_gce_board( 72 lsb_release_content=result.stdout) 73 except (error.AutoservRunError, error.AutoservSSHTimeout): 74 return False 75 76 def _modify_ssh_keys(self, to_add, to_remove): 77 """Modifies the list of ssh keys. 78 79 @param username: user name to add. 80 @param to_add: a list of new enties. 81 @param to_remove: a list of enties to be removed. 82 """ 83 keys = self.gce.GetCommonInstanceMetadata( 84 SSH_KEYS_METADATA_KEY) or '' 85 key_set = set(string.split(keys, '\n')) 86 new_key_set = (key_set | set(to_add)) - set(to_remove) 87 if key_set != new_key_set: 88 self.gce.SetCommonInstanceMetadata( 89 SSH_KEYS_METADATA_KEY, 90 string.join(list(new_key_set), '\n')) 91 92 def add_ssh_key(self, username, ssh_key): 93 """Adds a new SSH key in GCE metadata. 94 95 @param username: user name to add. 96 @param ssh_key: the key to add. 97 """ 98 self._modify_ssh_keys(['%s:%s' % (username, ssh_key)], []) 99 100 101 def del_ssh_key(self, username, ssh_key): 102 """Deletes the given SSH key from GCE metadata 103 104 @param username: user name to delete. 105 @param ssh_key: the key to delete. 106 """ 107 self._modify_ssh_keys([], ['%s:%s' % (username, ssh_key)]) 108 109 110 def get_release_version(self): 111 """Get the value of attribute CHROMEOS_RELEASE_VERSION from lsb-release. 112 113 @returns The version string in lsb-release, under attribute 114 CHROMEOS_RELEASE_VERSION. 115 """ 116 lsb_release_content = self.run( 117 'cat "%s"' % client_constants.LSB_RELEASE).stdout.strip() 118 return lsbrelease_utils.get_chromeos_release_version( 119 lsb_release_content=lsb_release_content) 120 121 def get_tmp_dir(self, parent=TMP_DIR): 122 """Return the pathname of a directory on the host suitable 123 for temporary file storage. 124 125 The directory and its content will be deleted automatically 126 on the destruction of the Host object that was used to obtain 127 it. 128 129 @param parent: Parent directory of the returned tmp dir. 130 131 @returns a path to the tmp directory on the host. 132 """ 133 if not parent.startswith(TMP_DIR): 134 parent = os.path.join(TMP_DIR, parent.lstrip(os.path.sep)) 135 self.run("mkdir -p %s" % parent) 136 template = os.path.join(parent, 'autoserv-XXXXXX') 137 dir_name = self.run_output("mktemp -d %s" % template) 138 self.tmp_dirs.append(dir_name) 139 return dir_name 140 141 142 def set_instance_metadata(self, key, value): 143 """Sets a single metadata value on the DUT instance. 144 145 @param key: Metadata key to be set. 146 @param value: New value, or None if the given key should be removed. 147 """ 148 self.gce.SetInstanceMetadata(self._gce_instance, key, value) 149 150 def stop(self): 151 """Stops the DUT instance 152 """ 153 self.gce.StopInstance(self._gce_instance) 154 155 def start(self): 156 """Starts the DUT instance 157 """ 158 self.gce.StartInstance(self._gce_instance) 159
