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 string 6import common 7 8from chromite.lib import gce 9 10from autotest_lib.client.common_lib import error 11from autotest_lib.client.common_lib import lsbrelease_utils 12from autotest_lib.client.cros import constants as client_constants 13from autotest_lib.server.hosts import abstract_ssh 14 15SSH_KEYS_METADATA_KEY = "sshKeys" 16 17def extract_arguments(args_dict): 18 """Extract GCE-specific arguments from arguments dictionary. 19 20 @param args_dict: dictionary of all arguments supplied to the test. 21 """ 22 23 return {k: v for k, v in args_dict.items() 24 if k in ('gce_project', 'gce_instance', 25 'gce_zone', 'gce_key_file')} 26 27 28class GceHost(abstract_ssh.AbstractSSHHost): 29 """GCE-specific subclass of Host.""" 30 31 def _initialize(self, hostname, gce_args=None, 32 *args, **dargs): 33 """Initializes this instance of GceHost. 34 35 @param hostname: the hostnname to be passed down to AbstractSSHHost. 36 @param gce_args: GCE-specific arguments extracted using 37 extract_arguments(). 38 """ 39 super(GceHost, self)._initialize(hostname=hostname, 40 *args, **dargs) 41 42 if gce_args: 43 self._gce_project = gce_args['gce_project'] 44 self._gce_zone = gce_args['gce_zone'] 45 self._gce_instance = gce_args['gce_instance'] 46 self._gce_key_file = gce_args['gce_key_file'] 47 else: 48 # TODO(andreyu): determine project, zone and instance names by 49 # querying metadata from the DUT instance 50 raise error.AutoservError('No GCE flags provided.') 51 52 self.gce = gce.GceContext.ForServiceAccountThreadSafe( 53 self._gce_project, self._gce_zone, self._gce_key_file) 54 55 56 def _modify_ssh_keys(self, to_add, to_remove): 57 """Modifies the list of ssh keys. 58 59 @param username: user name to add. 60 @param to_add: a list of new enties. 61 @param to_remove: a list of enties to be removed. 62 """ 63 keys = self.gce.GetCommonInstanceMetadata( 64 SSH_KEYS_METADATA_KEY) or '' 65 key_set = set(string.split(keys, '\n')) 66 new_key_set = (key_set | set(to_add)) - set(to_remove) 67 if key_set != new_key_set: 68 self.gce.SetCommonInstanceMetadata( 69 SSH_KEYS_METADATA_KEY, 70 string.join(list(new_key_set), '\n')) 71 72 def add_ssh_key(self, username, ssh_key): 73 """Adds a new SSH key in GCE metadata. 74 75 @param username: user name to add. 76 @param ssh_key: the key to add. 77 """ 78 self._modify_ssh_keys(['%s:%s' % (username, ssh_key)], []) 79 80 81 def del_ssh_key(self, username, ssh_key): 82 """Deletes the given SSH key from GCE metadata 83 84 @param username: user name to delete. 85 @param ssh_key: the key to delete. 86 """ 87 self._modify_ssh_keys([], ['%s:%s' % (username, ssh_key)]) 88 89 90 def get_release_version(self): 91 """Get the value of attribute CHROMEOS_RELEASE_VERSION from lsb-release. 92 93 @returns The version string in lsb-release, under attribute 94 CHROMEOS_RELEASE_VERSION. 95 """ 96 lsb_release_content = self.run( 97 'cat "%s"' % client_constants.LSB_RELEASE).stdout.strip() 98 return lsbrelease_utils.get_chromeos_release_version( 99 lsb_release_content=lsb_release_content) 100 101 102 def set_instance_metadata(self, key, value): 103 """Sets a single metadata value on the DUT instance. 104 105 @param key: Metadata key to be set. 106 @param value: New value, or None if the given key should be removed. 107 """ 108 self.gce.SetInstanceMetadata(self._gce_instance, key, value) 109 110 def stop(self): 111 """Stops the DUT instance 112 """ 113 self.gce.StopInstance(self._gce_instance) 114 115 def start(self): 116 """Starts the DUT instance 117 """ 118 self.gce.StartInstance(self._gce_instance) 119