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