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