1# Copyright 2016 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"""This class defines the Base Label classes."""
6
7
8import logging
9
10import common
11
12from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
13
14
15def forever_exists_decorate(exists):
16    """
17    Decorator for labels that should exist forever once applied.
18
19    We'll check if the label already exists on the host and return True if so.
20    Otherwise we'll check if the label should exist on the host.
21
22    @param exists: The exists method on the label class.
23    """
24    def exists_wrapper(self, host):
25        """
26        Wrapper around the label exists method.
27
28        @param self: The label object.
29        @param host: The host object to run methods on.
30
31        @returns True if the label already exists on the host, otherwise run
32            the exists method.
33        """
34        info = host.host_info_store.get()
35        return (self._NAME in info.labels) or exists(self, host)
36    return exists_wrapper
37
38
39class BaseLabel(object):
40    """
41    This class contains the scaffolding for the host-specific labels.
42
43    @property _NAME String that is either the label returned or a prefix of a
44                    generated label.
45    @property _LABEL_LIST List of label classes that this label generates its
46                          own labels from.  This class attribute is primarily
47                          for the LabelRetriever class to figure out what
48                          labels are generated from this label.  In most cases,
49                          the _NAME attribute gives us what we want, but in the
50                          special case where a label class is actually a
51                          collection of label classes, then this attribute
52                          comes into play.  For the example of
53                          testbed_label.ADBDeviceLabels, that class is really a
54                          collection of the adb devices' labels in that testbed
55                          so _NAME won't cut it.  Instead, we use _LABEL_LIST
56                          to tell LabelRetriever what list of label classes we
57                          are generating and thus are able to have a
58                          comprehensive list of the generated labels.
59    """
60
61    _NAME = None
62    _LABEL_LIST = []
63
64    def generate_labels(self, host):
65        """
66        Return the list of labels generated for the host.
67
68        @param host: The host object to check on.  Not needed here for base case
69                     but could be needed for subclasses.
70
71        @return a list of labels applicable to the host.
72        """
73        return [self._NAME]
74
75
76    def exists(self, host):
77        """
78        Checks the host if the label is applicable or not.
79
80        This method is geared for the type of labels that indicate if the host
81        has a feature (bluetooth, touchscreen, etc) and as such require
82        detection logic to determine if the label should be applicable to the
83        host or not.
84
85        @param host: The host object to check on.
86        """
87        raise NotImplementedError('exists not implemented')
88
89
90    def get(self, host):
91        """
92        Return the list of labels.
93
94        @param host: The host object to check on.
95        """
96        if self.exists(host):
97            return self.generate_labels(host)
98        else:
99            return []
100
101
102    def get_all_labels(self):
103        """
104        Return all possible labels generated by this label class.
105
106        @returns a tuple of sets, the first set is for labels that are prefixes
107            like 'os:android'.  The second set is for labels that are full
108            labels by themselves like 'bluetooth'.
109        """
110        # Another subclass takes care of prefixed labels so this is empty.
111        prefix_labels = set()
112        full_labels_list = (self._NAME if isinstance(self._NAME, list) else
113                            [self._NAME])
114        full_labels = set(full_labels_list)
115
116        return prefix_labels, full_labels
117
118
119class StringLabel(BaseLabel):
120    """
121    This class represents a string label that is dynamically generated.
122
123    This label class is used for the types of label that are always
124    present and will return at least one label out of a list of possible labels
125    (listed in _NAME).  It is required that the subclasses implement
126    generate_labels() since the label class will need to figure out which labels
127    to return.
128
129    _NAME must always be overridden by the subclass with all the possible
130    labels that this label detection class can return in order to allow for
131    accurate label updating.
132    """
133
134    def generate_labels(self, host):
135        raise NotImplementedError('generate_labels not implemented')
136
137
138    def exists(self, host):
139        """Set to true since it is assumed the label is always applicable."""
140        return True
141
142
143class StringPrefixLabel(StringLabel):
144    """
145    This class represents a string label that is dynamically generated.
146
147    This label class is used for the types of label that usually are always
148    present and indicate the os/board/etc type of the host.  The _NAME property
149    will be prepended with a colon to the generated labels like so:
150
151        _NAME = 'os'
152        generate_label() returns ['android']
153
154    The labels returned by this label class will be ['os:android'].
155    It is important that the _NAME attribute be overridden by the
156    subclass; otherwise, all labels returned will be prefixed with 'None:'.
157    """
158
159    def get(self, host):
160        """Return the list of labels with _NAME prefixed with a colon.
161
162        @param host: The host object to check on.
163        """
164        if self.exists(host):
165            return ['%s:%s' % (self._NAME, label)
166                    for label in self.generate_labels(host)]
167        else:
168            return []
169
170
171    def get_all_labels(self):
172        """
173        Return all possible labels generated by this label class.
174
175        @returns a tuple of sets, the first set is for labels that are prefixes
176            like 'os:android'.  The second set is for labels that are full
177            labels by themselves like 'bluetooth'.
178        """
179        # Since this is a prefix label class, we only care about
180        # prefixed_labels.  We'll need to append the ':' to the label name to
181        # make sure we only match on prefix labels.
182        full_labels = set()
183        prefix_labels = set(['%s:' % self._NAME])
184
185        return prefix_labels, full_labels
186
187
188class LabelRetriever(object):
189    """This class will assist in retrieving/updating the host labels."""
190
191    def _populate_known_labels(self, label_list):
192        """Create a list of known labels that is created through this class."""
193        for label_instance in label_list:
194            # If this instance has a list of label, recurse on that list.
195            if label_instance._LABEL_LIST:
196                self._populate_known_labels(label_instance._LABEL_LIST)
197                continue
198
199            prefixed_labels, full_labels = label_instance.get_all_labels()
200            self.label_prefix_names.update(prefixed_labels)
201            self.label_full_names.update(full_labels)
202
203
204    def __init__(self, label_list):
205        self._labels = label_list
206        # These two sets will contain the list of labels we can safely remove
207        # during the update_labels call.
208        self.label_full_names = set()
209        self.label_prefix_names = set()
210
211
212    def get_labels(self, host):
213        """
214        Retrieve the labels for the host.
215
216        @param host: The host to get the labels for.
217        """
218        labels = []
219        for label in self._labels:
220            logging.info('checking label %s', label.__class__.__name__)
221            try:
222                labels.extend(label.get(host))
223            except Exception:
224                logging.exception('error getting label %s.',
225                                  label.__class__.__name__)
226        return labels
227
228
229    def _is_known_label(self, label):
230        """
231        Checks if the label is a label known to the label detection framework.
232
233        We only delete labels that we might have created earlier.  There are
234        some labels we should not be removing (e.g. pool:bvt) that we
235        want to keep but won't be part of the new labels detected on the host.
236        To do that we compare the passed in label to our list of known labels
237        and if we get a match, we feel safe knowing we can remove the label.
238        Otherwise we leave that label alone since it was generated elsewhere.
239
240        @param label: The label to check if we want to skip or not.
241
242        @returns True to skip (which means to keep this label, False to remove.
243        """
244        return (label in self.label_full_names or
245                any([label.startswith(p) for p in self.label_prefix_names]))
246
247
248    def update_labels(self, host):
249        """
250        Retrieve the labels from the host and update if needed.
251
252        @param host: The host to update the labels for.
253        """
254        # If we haven't yet grabbed our list of known labels, do so now.
255        if not self.label_full_names and not self.label_prefix_names:
256            self._populate_known_labels(self._labels)
257
258        afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
259        old_labels = set(host._afe_host.labels)
260        logging.info('existing labels: %s', old_labels)
261        known_labels = set([l for l in old_labels
262                            if self._is_known_label(l)])
263        new_labels = set(self.get_labels(host))
264
265        # TODO(pprabhu) Replace this update logic using AfeHostInfoBackend.
266        # Remove old labels.
267        labels_to_remove = list(old_labels & (known_labels - new_labels))
268        if labels_to_remove:
269            logging.info('removing labels: %s', labels_to_remove)
270            afe.run('host_remove_labels', id=host.hostname,
271                    labels=labels_to_remove)
272
273        # Add in new labels that aren't already there.
274        labels_to_add = list(new_labels - old_labels)
275        if labels_to_add:
276            logging.info('adding labels: %s', labels_to_add)
277            afe.run('host_add_labels', id=host.hostname, labels=labels_to_add)
278