1# Copyright 2008 Google Inc. All Rights Reserved.
2
3"""
4The host module contains the objects and method used to
5manage a host in Autotest.
6
7The valid actions are:
8create:  adds host(s)
9delete:  deletes host(s)
10list:    lists host(s)
11stat:    displays host(s) information
12mod:     modifies host(s)
13jobs:    lists all jobs that ran on host(s)
14
15The common options are:
16-M|--mlist:   file containing a list of machines
17
18
19See topic_common.py for a High Level Design and Algorithm.
20
21"""
22import common
23import re
24import socket
25
26from autotest_lib.cli import action_common, rpc, topic_common
27from autotest_lib.client.bin import utils
28from autotest_lib.client.common_lib import error, host_protections
29from autotest_lib.server import frontend, hosts
30
31
32class host(topic_common.atest):
33    """Host class
34    atest host [create|delete|list|stat|mod|jobs] <options>"""
35    usage_action = '[create|delete|list|stat|mod|jobs]'
36    topic = msg_topic = 'host'
37    msg_items = '<hosts>'
38
39    protections = host_protections.Protection.names
40
41
42    def __init__(self):
43        """Add to the parser the options common to all the
44        host actions"""
45        super(host, self).__init__()
46
47        self.parser.add_option('-M', '--mlist',
48                               help='File listing the machines',
49                               type='string',
50                               default=None,
51                               metavar='MACHINE_FLIST')
52
53        self.topic_parse_info = topic_common.item_parse_info(
54            attribute_name='hosts',
55            filename_option='mlist',
56            use_leftover=True)
57
58
59    def _parse_lock_options(self, options):
60        if options.lock and options.unlock:
61            self.invalid_syntax('Only specify one of '
62                                '--lock and --unlock.')
63
64        if options.lock:
65            self.data['locked'] = True
66            self.messages.append('Locked host')
67        elif options.unlock:
68            self.data['locked'] = False
69            self.data['lock_reason'] = ''
70            self.messages.append('Unlocked host')
71
72        if options.lock and options.lock_reason:
73            self.data['lock_reason'] = options.lock_reason
74
75
76    def _cleanup_labels(self, labels, platform=None):
77        """Removes the platform label from the overall labels"""
78        if platform:
79            return [label for label in labels
80                    if label != platform]
81        else:
82            try:
83                return [label for label in labels
84                        if not label['platform']]
85            except TypeError:
86                # This is a hack - the server will soon
87                # do this, so all this code should be removed.
88                return labels
89
90
91    def get_items(self):
92        return self.hosts
93
94
95class host_help(host):
96    """Just here to get the atest logic working.
97    Usage is set by its parent"""
98    pass
99
100
101class host_list(action_common.atest_list, host):
102    """atest host list [--mlist <file>|<hosts>] [--label <label>]
103       [--status <status1,status2>] [--acl <ACL>] [--user <user>]"""
104
105    def __init__(self):
106        super(host_list, self).__init__()
107
108        self.parser.add_option('-b', '--label',
109                               default='',
110                               help='Only list hosts with all these labels '
111                               '(comma separated)')
112        self.parser.add_option('-s', '--status',
113                               default='',
114                               help='Only list hosts with any of these '
115                               'statuses (comma separated)')
116        self.parser.add_option('-a', '--acl',
117                               default='',
118                               help='Only list hosts within this ACL')
119        self.parser.add_option('-u', '--user',
120                               default='',
121                               help='Only list hosts available to this user')
122        self.parser.add_option('-N', '--hostnames-only', help='Only return '
123                               'hostnames for the machines queried.',
124                               action='store_true')
125        self.parser.add_option('--locked',
126                               default=False,
127                               help='Only list locked hosts',
128                               action='store_true')
129        self.parser.add_option('--unlocked',
130                               default=False,
131                               help='Only list unlocked hosts',
132                               action='store_true')
133
134
135
136    def parse(self):
137        """Consume the specific options"""
138        label_info = topic_common.item_parse_info(attribute_name='labels',
139                                                  inline_option='label')
140
141        (options, leftover) = super(host_list, self).parse([label_info])
142
143        self.status = options.status
144        self.acl = options.acl
145        self.user = options.user
146        self.hostnames_only = options.hostnames_only
147
148        if options.locked and options.unlocked:
149            self.invalid_syntax('--locked and --unlocked are '
150                                'mutually exclusive')
151        self.locked = options.locked
152        self.unlocked = options.unlocked
153        return (options, leftover)
154
155
156    def execute(self):
157        """Execute 'atest host list'."""
158        filters = {}
159        check_results = {}
160        if self.hosts:
161            filters['hostname__in'] = self.hosts
162            check_results['hostname__in'] = 'hostname'
163
164        if self.labels:
165            if len(self.labels) == 1:
166                # This is needed for labels with wildcards (x86*)
167                filters['labels__name__in'] = self.labels
168                check_results['labels__name__in'] = None
169            else:
170                filters['multiple_labels'] = self.labels
171                check_results['multiple_labels'] = None
172
173        if self.status:
174            statuses = self.status.split(',')
175            statuses = [status.strip() for status in statuses
176                        if status.strip()]
177
178            filters['status__in'] = statuses
179            check_results['status__in'] = None
180
181        if self.acl:
182            filters['aclgroup__name'] = self.acl
183            check_results['aclgroup__name'] = None
184        if self.user:
185            filters['aclgroup__users__login'] = self.user
186            check_results['aclgroup__users__login'] = None
187
188        if self.locked or self.unlocked:
189            filters['locked'] = self.locked
190            check_results['locked'] = None
191
192        return super(host_list, self).execute(op='get_hosts',
193                                              filters=filters,
194                                              check_results=check_results)
195
196
197    def output(self, results):
198        """Print output of 'atest host list'.
199
200        @param results: the results to be printed.
201        """
202        if results:
203            # Remove the platform from the labels.
204            for result in results:
205                result['labels'] = self._cleanup_labels(result['labels'],
206                                                        result['platform'])
207        if self.hostnames_only:
208            self.print_list(results, key='hostname')
209        else:
210            keys = ['hostname', 'status',
211                    'shard', 'locked', 'lock_reason', 'locked_by', 'platform',
212                    'labels']
213            super(host_list, self).output(results, keys=keys)
214
215
216class host_stat(host):
217    """atest host stat --mlist <file>|<hosts>"""
218    usage_action = 'stat'
219
220    def execute(self):
221        """Execute 'atest host stat'."""
222        results = []
223        # Convert wildcards into real host stats.
224        existing_hosts = []
225        for host in self.hosts:
226            if host.endswith('*'):
227                stats = self.execute_rpc('get_hosts',
228                                         hostname__startswith=host.rstrip('*'))
229                if len(stats) == 0:
230                    self.failure('No hosts matching %s' % host, item=host,
231                                 what_failed='Failed to stat')
232                    continue
233            else:
234                stats = self.execute_rpc('get_hosts', hostname=host)
235                if len(stats) == 0:
236                    self.failure('Unknown host %s' % host, item=host,
237                                 what_failed='Failed to stat')
238                    continue
239            existing_hosts.extend(stats)
240
241        for stat in existing_hosts:
242            host = stat['hostname']
243            # The host exists, these should succeed
244            acls = self.execute_rpc('get_acl_groups', hosts__hostname=host)
245
246            labels = self.execute_rpc('get_labels', host__hostname=host)
247            results.append([[stat], acls, labels, stat['attributes']])
248        return results
249
250
251    def output(self, results):
252        """Print output of 'atest host stat'.
253
254        @param results: the results to be printed.
255        """
256        for stats, acls, labels, attributes in results:
257            print '-'*5
258            self.print_fields(stats,
259                              keys=['hostname', 'platform',
260                                    'status', 'locked', 'locked_by',
261                                    'lock_time', 'lock_reason', 'protection',])
262            self.print_by_ids(acls, 'ACLs', line_before=True)
263            labels = self._cleanup_labels(labels)
264            self.print_by_ids(labels, 'Labels', line_before=True)
265            self.print_dict(attributes, 'Host Attributes', line_before=True)
266
267
268class host_jobs(host):
269    """atest host jobs [--max-query] --mlist <file>|<hosts>"""
270    usage_action = 'jobs'
271
272    def __init__(self):
273        super(host_jobs, self).__init__()
274        self.parser.add_option('-q', '--max-query',
275                               help='Limits the number of results '
276                               '(20 by default)',
277                               type='int', default=20)
278
279
280    def parse(self):
281        """Consume the specific options"""
282        (options, leftover) = super(host_jobs, self).parse()
283        self.max_queries = options.max_query
284        return (options, leftover)
285
286
287    def execute(self):
288        """Execute 'atest host jobs'."""
289        results = []
290        real_hosts = []
291        for host in self.hosts:
292            if host.endswith('*'):
293                stats = self.execute_rpc('get_hosts',
294                                         hostname__startswith=host.rstrip('*'))
295                if len(stats) == 0:
296                    self.failure('No host matching %s' % host, item=host,
297                                 what_failed='Failed to stat')
298                [real_hosts.append(stat['hostname']) for stat in stats]
299            else:
300                real_hosts.append(host)
301
302        for host in real_hosts:
303            queue_entries = self.execute_rpc('get_host_queue_entries',
304                                             host__hostname=host,
305                                             query_limit=self.max_queries,
306                                             sort_by=['-job__id'])
307            jobs = []
308            for entry in queue_entries:
309                job = {'job_id': entry['job']['id'],
310                       'job_owner': entry['job']['owner'],
311                       'job_name': entry['job']['name'],
312                       'status': entry['status']}
313                jobs.append(job)
314            results.append((host, jobs))
315        return results
316
317
318    def output(self, results):
319        """Print output of 'atest host jobs'.
320
321        @param results: the results to be printed.
322        """
323        for host, jobs in results:
324            print '-'*5
325            print 'Hostname: %s' % host
326            self.print_table(jobs, keys_header=['job_id',
327                                                'job_owner',
328                                                'job_name',
329                                                'status'])
330
331class BaseHostModCreate(host):
332    """The base class for host_mod and host_create"""
333    # Matches one attribute=value pair
334    attribute_regex = r'(?P<attribute>\w+)=(?P<value>.+)?'
335
336    def __init__(self):
337        """Add the options shared between host mod and host create actions."""
338        self.messages = []
339        self.host_ids = {}
340        super(BaseHostModCreate, self).__init__()
341        self.parser.add_option('-l', '--lock',
342                               help='Lock hosts',
343                               action='store_true')
344        self.parser.add_option('-u', '--unlock',
345                               help='Unlock hosts',
346                               action='store_true')
347        self.parser.add_option('-r', '--lock_reason',
348                               help='Reason for locking hosts',
349                               default='')
350        self.parser.add_option('-p', '--protection', type='choice',
351                               help=('Set the protection level on a host.  '
352                                     'Must be one of: %s' %
353                                     ', '.join('"%s"' % p
354                                               for p in self.protections)),
355                               choices=self.protections)
356        self._attributes = []
357        self.parser.add_option('--attribute', '-i',
358                               help=('Host attribute to add or change. Format '
359                                     'is <attribute>=<value>. Multiple '
360                                     'attributes can be set by passing the '
361                                     'argument multiple times. Attributes can '
362                                     'be unset by providing an empty value.'),
363                               action='append')
364        self.parser.add_option('-b', '--labels',
365                               help='Comma separated list of labels')
366        self.parser.add_option('-B', '--blist',
367                               help='File listing the labels',
368                               type='string',
369                               metavar='LABEL_FLIST')
370        self.parser.add_option('-a', '--acls',
371                               help='Comma separated list of ACLs')
372        self.parser.add_option('-A', '--alist',
373                               help='File listing the acls',
374                               type='string',
375                               metavar='ACL_FLIST')
376        self.parser.add_option('-t', '--platform',
377                               help='Sets the platform label')
378
379
380    def parse(self):
381        """Consume the options common to host create and host mod.
382        """
383        label_info = topic_common.item_parse_info(attribute_name='labels',
384                                                 inline_option='labels',
385                                                 filename_option='blist')
386        acl_info = topic_common.item_parse_info(attribute_name='acls',
387                                                inline_option='acls',
388                                                filename_option='alist')
389
390        (options, leftover) = super(BaseHostModCreate, self).parse([label_info,
391                                                              acl_info],
392                                                             req_items='hosts')
393
394        self._parse_lock_options(options)
395
396        if options.protection:
397            self.data['protection'] = options.protection
398            self.messages.append('Protection set to "%s"' % options.protection)
399
400        self.attributes = {}
401        if options.attribute:
402            for pair in options.attribute:
403                m = re.match(self.attribute_regex, pair)
404                if not m:
405                    raise topic_common.CliError('Attribute must be in key=value '
406                                                'syntax.')
407                elif m.group('attribute') in self.attributes:
408                    raise topic_common.CliError(
409                            'Multiple values provided for attribute '
410                            '%s.' % m.group('attribute'))
411                self.attributes[m.group('attribute')] = m.group('value')
412
413        self.platform = options.platform
414        return (options, leftover)
415
416
417    def _set_acls(self, hosts, acls):
418        """Add hosts to acls (and remove from all other acls).
419
420        @param hosts: list of hostnames
421        @param acls: list of acl names
422        """
423        # Remove from all ACLs except 'Everyone' and ACLs in list
424        # Skip hosts that don't exist
425        for host in hosts:
426            if host not in self.host_ids:
427                continue
428            host_id = self.host_ids[host]
429            for a in self.execute_rpc('get_acl_groups', hosts=host_id):
430                if a['name'] not in self.acls and a['id'] != 1:
431                    self.execute_rpc('acl_group_remove_hosts', id=a['id'],
432                                     hosts=self.hosts)
433
434        # Add hosts to the ACLs
435        self.check_and_create_items('get_acl_groups', 'add_acl_group',
436                                    self.acls)
437        for a in acls:
438            self.execute_rpc('acl_group_add_hosts', id=a, hosts=hosts)
439
440
441    def _remove_labels(self, host, condition):
442        """Remove all labels from host that meet condition(label).
443
444        @param host: hostname
445        @param condition: callable that returns bool when given a label
446        """
447        if host in self.host_ids:
448            host_id = self.host_ids[host]
449            labels_to_remove = []
450            for l in self.execute_rpc('get_labels', host=host_id):
451                if condition(l):
452                    labels_to_remove.append(l['id'])
453            if labels_to_remove:
454                self.execute_rpc('host_remove_labels', id=host_id,
455                                 labels=labels_to_remove)
456
457
458    def _set_labels(self, host, labels):
459        """Apply labels to host (and remove all other labels).
460
461        @param host: hostname
462        @param labels: list of label names
463        """
464        condition = lambda l: l['name'] not in labels and not l['platform']
465        self._remove_labels(host, condition)
466        self.check_and_create_items('get_labels', 'add_label', labels)
467        self.execute_rpc('host_add_labels', id=host, labels=labels)
468
469
470    def _set_platform_label(self, host, platform_label):
471        """Apply the platform label to host (and remove existing).
472
473        @param host: hostname
474        @param platform_label: platform label's name
475        """
476        self._remove_labels(host, lambda l: l['platform'])
477        self.check_and_create_items('get_labels', 'add_label', [platform_label],
478                                    platform=True)
479        self.execute_rpc('host_add_labels', id=host, labels=[platform_label])
480
481
482    def _set_attributes(self, host, attributes):
483        """Set attributes on host.
484
485        @param host: hostname
486        @param attributes: attribute dictionary
487        """
488        for attr, value in self.attributes.iteritems():
489            self.execute_rpc('set_host_attribute', attribute=attr,
490                             value=value, hostname=host)
491
492
493class host_mod(BaseHostModCreate):
494    """atest host mod [--lock|--unlock --force_modify_locking
495    --platform <arch>
496    --labels <labels>|--blist <label_file>
497    --acls <acls>|--alist <acl_file>
498    --protection <protection_type>
499    --attributes <attr>=<value>;<attr>=<value>
500    --mlist <mach_file>] <hosts>"""
501    usage_action = 'mod'
502
503    def __init__(self):
504        """Add the options specific to the mod action"""
505        super(host_mod, self).__init__()
506        self.parser.add_option('-f', '--force_modify_locking',
507                               help='Forcefully lock\unlock a host',
508                               action='store_true')
509        self.parser.add_option('--remove_acls',
510                               help='Remove all active acls.',
511                               action='store_true')
512        self.parser.add_option('--remove_labels',
513                               help='Remove all labels.',
514                               action='store_true')
515
516
517    def parse(self):
518        """Consume the specific options"""
519        (options, leftover) = super(host_mod, self).parse()
520
521        if options.force_modify_locking:
522             self.data['force_modify_locking'] = True
523
524        self.remove_acls = options.remove_acls
525        self.remove_labels = options.remove_labels
526
527        return (options, leftover)
528
529
530    def execute(self):
531        """Execute 'atest host mod'."""
532        successes = []
533        for host in self.execute_rpc('get_hosts', hostname__in=self.hosts):
534            self.host_ids[host['hostname']] = host['id']
535        for host in self.hosts:
536            if host not in self.host_ids:
537                self.failure('Cannot modify non-existant host %s.' % host)
538                continue
539            host_id = self.host_ids[host]
540
541            try:
542                if self.data:
543                    self.execute_rpc('modify_host', item=host,
544                                     id=host, **self.data)
545
546                if self.attributes:
547                    self._set_attributes(host, self.attributes)
548
549                if self.labels or self.remove_labels:
550                    self._set_labels(host, self.labels)
551
552                if self.platform:
553                    self._set_platform_label(host, self.platform)
554
555                # TODO: Make the AFE return True or False,
556                # especially for lock
557                successes.append(host)
558            except topic_common.CliError, full_error:
559                # Already logged by execute_rpc()
560                pass
561
562        if self.acls or self.remove_acls:
563            self._set_acls(self.hosts, self.acls)
564
565        return successes
566
567
568    def output(self, hosts):
569        """Print output of 'atest host mod'.
570
571        @param hosts: the host list to be printed.
572        """
573        for msg in self.messages:
574            self.print_wrapped(msg, hosts)
575
576
577class HostInfo(object):
578    """Store host information so we don't have to keep looking it up."""
579    def __init__(self, hostname, platform, labels):
580        self.hostname = hostname
581        self.platform = platform
582        self.labels = labels
583
584
585class host_create(BaseHostModCreate):
586    """atest host create [--lock|--unlock --platform <arch>
587    --labels <labels>|--blist <label_file>
588    --acls <acls>|--alist <acl_file>
589    --protection <protection_type>
590    --attributes <attr>=<value>;<attr>=<value>
591    --mlist <mach_file>] <hosts>"""
592    usage_action = 'create'
593
594    def parse(self):
595        """Option logic specific to create action.
596        """
597        (options, leftovers) = super(host_create, self).parse()
598        self.locked = options.lock
599        if 'serials' in self.attributes:
600            if len(self.hosts) > 1:
601                raise topic_common.CliError('Can not specify serials with '
602                                            'multiple hosts.')
603
604
605    @classmethod
606    def construct_without_parse(
607            cls, web_server, hosts, platform=None,
608            locked=False, lock_reason='', labels=[], acls=[],
609            protection=host_protections.Protection.NO_PROTECTION):
610        """Construct a host_create object and fill in data from args.
611
612        Do not need to call parse after the construction.
613
614        Return an object of site_host_create ready to execute.
615
616        @param web_server: A string specifies the autotest webserver url.
617            It is needed to setup comm to make rpc.
618        @param hosts: A list of hostnames as strings.
619        @param platform: A string or None.
620        @param locked: A boolean.
621        @param lock_reason: A string.
622        @param labels: A list of labels as strings.
623        @param acls: A list of acls as strings.
624        @param protection: An enum defined in host_protections.
625        """
626        obj = cls()
627        obj.web_server = web_server
628        try:
629            # Setup stuff needed for afe comm.
630            obj.afe = rpc.afe_comm(web_server)
631        except rpc.AuthError, s:
632            obj.failure(str(s), fatal=True)
633        obj.hosts = hosts
634        obj.platform = platform
635        obj.locked = locked
636        if locked and lock_reason.strip():
637            obj.data['lock_reason'] = lock_reason.strip()
638        obj.labels = labels
639        obj.acls = acls
640        if protection:
641            obj.data['protection'] = protection
642        obj.attributes = {}
643        return obj
644
645
646    def _detect_host_info(self, host):
647        """Detect platform and labels from the host.
648
649        @param host: hostname
650
651        @return: HostInfo object
652        """
653        # Mock an afe_host object so that the host is constructed as if the
654        # data was already in afe
655        data = {'attributes': self.attributes, 'labels': self.labels}
656        afe_host = frontend.Host(None, data)
657        machine = {'hostname': host, 'afe_host': afe_host}
658        try:
659            if utils.ping(host, tries=1, deadline=1) == 0:
660                serials = self.attributes.get('serials', '').split(',')
661                if serials and len(serials) > 1:
662                    host_dut = hosts.create_testbed(machine,
663                                                    adb_serials=serials)
664                else:
665                    adb_serial = self.attributes.get('serials')
666                    host_dut = hosts.create_host(machine,
667                                                 adb_serial=adb_serial)
668
669                host_info = HostInfo(host, host_dut.get_platform(),
670                                     host_dut.get_labels())
671                # Clean host to make sure nothing left after calling it,
672                # e.g. tunnels.
673                if hasattr(host_dut, 'close'):
674                    host_dut.close()
675            else:
676                # Can't ping the host, use default information.
677                host_info = HostInfo(host, None, [])
678        except (socket.gaierror, error.AutoservRunError,
679                error.AutoservSSHTimeout):
680            # We may be adding a host that does not exist yet or we can't
681            # reach due to hostname/address issues or if the host is down.
682            host_info = HostInfo(host, None, [])
683        return host_info
684
685
686    def _execute_add_one_host(self, host):
687        # Always add the hosts as locked to avoid the host
688        # being picked up by the scheduler before it's ACL'ed.
689        self.data['locked'] = True
690        if not self.locked:
691            self.data['lock_reason'] = 'Forced lock on device creation'
692        self.execute_rpc('add_host', hostname=host, status="Ready", **self.data)
693
694        # If there are labels avaliable for host, use them.
695        host_info = self._detect_host_info(host)
696        labels = set(self.labels)
697        if host_info.labels:
698            labels.update(host_info.labels)
699
700        if labels:
701            self._set_labels(host, list(labels))
702
703        # Now add the platform label.
704        # If a platform was not provided and we were able to retrieve it
705        # from the host, use the retrieved platform.
706        platform = self.platform if self.platform else host_info.platform
707        if platform:
708            self._set_platform_label(host, platform)
709
710        if self.attributes:
711            self._set_attributes(host, self.attributes)
712
713
714    def execute(self):
715        """Execute 'atest host create'."""
716        successful_hosts = []
717        for host in self.hosts:
718            try:
719                self._execute_add_one_host(host)
720                successful_hosts.append(host)
721            except topic_common.CliError:
722                pass
723
724        if successful_hosts:
725            self._set_acls(successful_hosts, self.acls)
726
727            if not self.locked:
728                for host in successful_hosts:
729                    self.execute_rpc('modify_host', id=host, locked=False,
730                                     lock_reason='')
731        return successful_hosts
732
733
734    def output(self, hosts):
735        """Print output of 'atest host create'.
736
737        @param hosts: the added host list to be printed.
738        """
739        self.print_wrapped('Added host', hosts)
740
741
742class host_delete(action_common.atest_delete, host):
743    """atest host delete [--mlist <mach_file>] <hosts>"""
744    pass
745