1#!/usr/bin/python3 2 3# Copyright (c) 2013 The Chromium OS Authors. All rights reserved. 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6"""A small wrapper script, iterates through 7the known hosts and tries to call get_labels() 8to discover host functionality, and adds these 9detected labels to host. 10 11Limitations: 12 - Does not keep a count of how many labels were 13 actually added. 14 - If a label is added by this script because it 15 is detected as supported by get_labels, but later becomes 16 unsupported, this script has no way to know that it 17 should be removed, so it will remain attached to the host. 18 See crosbug.com/38569 19""" 20 21 22from __future__ import absolute_import 23from __future__ import division 24from __future__ import print_function 25 26from multiprocessing import pool 27import logging 28import socket 29import argparse 30import sys 31 32import common 33 34from autotest_lib.server import hosts 35from autotest_lib.server import frontend 36from autotest_lib.client.common_lib import error 37 38 39# A list of label prefix that each dut should only have one of such label with 40# the given prefix, e.g., a dut can't have both labels of power:battery and 41# power:AC_only. 42SINGLETON_LABEL_PREFIX = ['power:'] 43 44def add_missing_labels(afe, hostname): 45 """ 46 Queries the detectable labels supported by the given host, 47 and adds those labels to the host. 48 49 @param afe: A frontend.AFE() instance. 50 @param hostname: The host to query and update. 51 52 @return: True on success. 53 False on failure to fetch labels or to add any individual label. 54 """ 55 host = None 56 try: 57 host = hosts.create_host(hostname) 58 labels = host.get_labels() 59 except socket.gaierror: 60 logging.warning('Unable to establish ssh connection to hostname ' 61 '%s. Skipping.', hostname) 62 return False 63 except error.AutoservError: 64 logging.warning('Unable to query labels on hostname %s. Skipping.', 65 hostname) 66 return False 67 finally: 68 if host: 69 host.close() 70 71 afe_host = afe.get_hosts(hostname=hostname)[0] 72 73 label_matches = afe.get_labels(name__in=labels) 74 75 for label in label_matches: 76 singleton_prefixes = [p for p in SINGLETON_LABEL_PREFIX 77 if label.name.startswith(p)] 78 if len(singleton_prefixes) == 1: 79 singleton_prefix = singleton_prefixes[0] 80 # Delete existing label with `singleton_prefix` 81 labels_to_delete = [l for l in afe_host.labels 82 if l.startswith(singleton_prefix) and 83 not l in labels] 84 if labels_to_delete: 85 logging.warning('Removing label %s', labels_to_delete) 86 afe_labels_to_delete = afe.get_labels(name__in=labels_to_delete) 87 for afe_label in afe_labels_to_delete: 88 afe_label.remove_hosts(hosts=[hostname]) 89 label.add_hosts(hosts=[hostname]) 90 91 missing_labels = set(labels) - set([l.name for l in label_matches]) 92 93 if missing_labels: 94 for label in missing_labels: 95 logging.warning('Unable to add label %s to host %s. ' 96 'Skipping unknown label.', label, hostname) 97 return False 98 99 return True 100 101 102def main(): 103 """" 104 Entry point for add_detected_host_labels script. 105 """ 106 107 parser = argparse.ArgumentParser() 108 parser.add_argument('-s', '--silent', dest='silent', action='store_true', 109 help='Suppress all but critical logging messages.') 110 parser.add_argument('-i', '--info', dest='info_only', action='store_true', 111 help='Suppress logging messages below INFO priority.') 112 parser.add_argument('-m', '--machines', dest='machines', 113 help='Comma separated list of machines to check.') 114 options = parser.parse_args() 115 116 if options.silent and options.info_only: 117 print('The -i and -s flags cannot be used together.') 118 parser.print_help() 119 return 0 120 121 122 if options.silent: 123 logging.disable(logging.CRITICAL) 124 125 if options.info_only: 126 logging.disable(logging.DEBUG) 127 128 threadpool = pool.ThreadPool() 129 afe = frontend.AFE() 130 131 if options.machines: 132 hostnames = [m.strip() for m in options.machines.split(',')] 133 else: 134 hostnames = afe.get_hostnames() 135 successes = sum(threadpool.imap_unordered( 136 lambda x: add_missing_labels(afe, x), 137 hostnames)) 138 attempts = len(hostnames) 139 140 logging.info('Label updating finished. Failed update on %d out of %d ' 141 'hosts.', attempts-successes, attempts) 142 143 return 0 144 145 146if __name__ == '__main__': 147 sys.exit(main()) 148