• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 random
24import re
25import socket
26
27from autotest_lib.cli import action_common, rpc, topic_common, skylab_utils
28from autotest_lib.client.bin import utils as bin_utils
29from autotest_lib.client.common_lib import error, host_protections
30from autotest_lib.server import frontend, hosts
31from autotest_lib.server.hosts import host_info
32
33
34try:
35    from skylab_inventory import text_manager
36    from skylab_inventory.lib import device
37    from skylab_inventory.lib import server as skylab_server
38except ImportError:
39    pass
40
41
42MIGRATED_HOST_SUFFIX = '-migrated-do-not-use'
43
44
45class host(topic_common.atest):
46    """Host class
47    atest host [create|delete|list|stat|mod|jobs|rename|migrate] <options>"""
48    usage_action = '[create|delete|list|stat|mod|jobs|rename|migrate]'
49    topic = msg_topic = 'host'
50    msg_items = '<hosts>'
51
52    protections = host_protections.Protection.names
53
54
55    def __init__(self):
56        """Add to the parser the options common to all the
57        host actions"""
58        super(host, self).__init__()
59
60        self.parser.add_option('-M', '--mlist',
61                               help='File listing the machines',
62                               type='string',
63                               default=None,
64                               metavar='MACHINE_FLIST')
65
66        self.topic_parse_info = topic_common.item_parse_info(
67            attribute_name='hosts',
68            filename_option='mlist',
69            use_leftover=True)
70
71
72    def _parse_lock_options(self, options):
73        if options.lock and options.unlock:
74            self.invalid_syntax('Only specify one of '
75                                '--lock and --unlock.')
76
77        self.lock = options.lock
78        self.unlock = options.unlock
79        self.lock_reason = options.lock_reason
80
81        if options.lock:
82            self.data['locked'] = True
83            self.messages.append('Locked host')
84        elif options.unlock:
85            self.data['locked'] = False
86            self.data['lock_reason'] = ''
87            self.messages.append('Unlocked host')
88
89        if options.lock and options.lock_reason:
90            self.data['lock_reason'] = options.lock_reason
91
92
93    def _cleanup_labels(self, labels, platform=None):
94        """Removes the platform label from the overall labels"""
95        if platform:
96            return [label for label in labels
97                    if label != platform]
98        else:
99            try:
100                return [label for label in labels
101                        if not label['platform']]
102            except TypeError:
103                # This is a hack - the server will soon
104                # do this, so all this code should be removed.
105                return labels
106
107
108    def get_items(self):
109        return self.hosts
110
111
112class host_help(host):
113    """Just here to get the atest logic working.
114    Usage is set by its parent"""
115    pass
116
117
118class host_list(action_common.atest_list, host):
119    """atest host list [--mlist <file>|<hosts>] [--label <label>]
120       [--status <status1,status2>] [--acl <ACL>] [--user <user>]"""
121
122    def __init__(self):
123        super(host_list, self).__init__()
124
125        self.parser.add_option('-b', '--label',
126                               default='',
127                               help='Only list hosts with all these labels '
128                               '(comma separated). When --skylab is provided, '
129                               'a label must be in the format of '
130                               'label-key:label-value (e.g., board:lumpy).')
131        self.parser.add_option('-s', '--status',
132                               default='',
133                               help='Only list hosts with any of these '
134                               'statuses (comma separated)')
135        self.parser.add_option('-a', '--acl',
136                               default='',
137                               help=('Only list hosts within this ACL. %s' %
138                                     skylab_utils.MSG_INVALID_IN_SKYLAB))
139        self.parser.add_option('-u', '--user',
140                               default='',
141                               help=('Only list hosts available to this user. '
142                                     '%s' % skylab_utils.MSG_INVALID_IN_SKYLAB))
143        self.parser.add_option('-N', '--hostnames-only', help='Only return '
144                               'hostnames for the machines queried.',
145                               action='store_true')
146        self.parser.add_option('--locked',
147                               default=False,
148                               help='Only list locked hosts',
149                               action='store_true')
150        self.parser.add_option('--unlocked',
151                               default=False,
152                               help='Only list unlocked hosts',
153                               action='store_true')
154        self.parser.add_option('--full-output',
155                               default=False,
156                               help=('Print out the full content of the hosts. '
157                                     'Only supported with --skylab.'),
158                               action='store_true',
159                               dest='full_output')
160
161        self.add_skylab_options()
162
163
164    def parse(self):
165        """Consume the specific options"""
166        label_info = topic_common.item_parse_info(attribute_name='labels',
167                                                  inline_option='label')
168
169        (options, leftover) = super(host_list, self).parse([label_info])
170
171        self.status = options.status
172        self.acl = options.acl
173        self.user = options.user
174        self.hostnames_only = options.hostnames_only
175
176        if options.locked and options.unlocked:
177            self.invalid_syntax('--locked and --unlocked are '
178                                'mutually exclusive')
179
180        self.locked = options.locked
181        self.unlocked = options.unlocked
182        self.label_map = None
183
184        if self.skylab:
185            if options.user or options.acl or options.status:
186                self.invalid_syntax('--user, --acl or --status is not '
187                                    'supported with --skylab.')
188            self.full_output = options.full_output
189            if self.full_output and self.hostnames_only:
190                self.invalid_syntax('--full-output is conflicted with '
191                                    '--hostnames-only.')
192
193            if self.labels:
194                self.label_map = device.convert_to_label_map(self.labels)
195        else:
196            if options.full_output:
197                self.invalid_syntax('--full_output is only supported with '
198                                    '--skylab.')
199
200        return (options, leftover)
201
202
203    def execute_skylab(self):
204        """Execute 'atest host list' with --skylab."""
205        inventory_repo = skylab_utils.InventoryRepo(self.inventory_repo_dir)
206        inventory_repo.initialize()
207        lab = text_manager.load_lab(inventory_repo.get_data_dir())
208
209        # TODO(nxia): support filtering on run-time labels and status.
210        return device.get_devices(
211            lab,
212            'duts',
213            self.environment,
214            label_map=self.label_map,
215            hostnames=self.hosts,
216            locked=self.locked,
217            unlocked=self.unlocked)
218
219
220    def execute(self):
221        """Execute 'atest host list'."""
222        if self.skylab:
223            return self.execute_skylab()
224
225        filters = {}
226        check_results = {}
227        if self.hosts:
228            filters['hostname__in'] = self.hosts
229            check_results['hostname__in'] = 'hostname'
230
231        if self.labels:
232            if len(self.labels) == 1:
233                # This is needed for labels with wildcards (x86*)
234                filters['labels__name__in'] = self.labels
235                check_results['labels__name__in'] = None
236            else:
237                filters['multiple_labels'] = self.labels
238                check_results['multiple_labels'] = None
239
240        if self.status:
241            statuses = self.status.split(',')
242            statuses = [status.strip() for status in statuses
243                        if status.strip()]
244
245            filters['status__in'] = statuses
246            check_results['status__in'] = None
247
248        if self.acl:
249            filters['aclgroup__name'] = self.acl
250            check_results['aclgroup__name'] = None
251        if self.user:
252            filters['aclgroup__users__login'] = self.user
253            check_results['aclgroup__users__login'] = None
254
255        if self.locked or self.unlocked:
256            filters['locked'] = self.locked
257            check_results['locked'] = None
258
259        return super(host_list, self).execute(op='get_hosts',
260                                              filters=filters,
261                                              check_results=check_results)
262
263
264    def output(self, results):
265        """Print output of 'atest host list'.
266
267        @param results: the results to be printed.
268        """
269        if results and not self.skylab:
270            # Remove the platform from the labels.
271            for result in results:
272                result['labels'] = self._cleanup_labels(result['labels'],
273                                                        result['platform'])
274        if self.skylab and self.full_output:
275            print results
276            return
277
278        if self.skylab:
279            results = device.convert_to_autotest_hosts(results)
280
281        if self.hostnames_only:
282            self.print_list(results, key='hostname')
283        else:
284            keys = ['hostname', 'status', 'shard', 'locked', 'lock_reason',
285                    'locked_by', 'platform', 'labels']
286            super(host_list, self).output(results, keys=keys)
287
288
289class host_stat(host):
290    """atest host stat --mlist <file>|<hosts>"""
291    usage_action = 'stat'
292
293    def execute(self):
294        """Execute 'atest host stat'."""
295        results = []
296        # Convert wildcards into real host stats.
297        existing_hosts = []
298        for host in self.hosts:
299            if host.endswith('*'):
300                stats = self.execute_rpc('get_hosts',
301                                         hostname__startswith=host.rstrip('*'))
302                if len(stats) == 0:
303                    self.failure('No hosts matching %s' % host, item=host,
304                                 what_failed='Failed to stat')
305                    continue
306            else:
307                stats = self.execute_rpc('get_hosts', hostname=host)
308                if len(stats) == 0:
309                    self.failure('Unknown host %s' % host, item=host,
310                                 what_failed='Failed to stat')
311                    continue
312            existing_hosts.extend(stats)
313
314        for stat in existing_hosts:
315            host = stat['hostname']
316            # The host exists, these should succeed
317            acls = self.execute_rpc('get_acl_groups', hosts__hostname=host)
318
319            labels = self.execute_rpc('get_labels', host__hostname=host)
320            results.append([[stat], acls, labels, stat['attributes']])
321        return results
322
323
324    def output(self, results):
325        """Print output of 'atest host stat'.
326
327        @param results: the results to be printed.
328        """
329        for stats, acls, labels, attributes in results:
330            print '-'*5
331            self.print_fields(stats,
332                              keys=['hostname', 'id', 'platform',
333                                    'status', 'locked', 'locked_by',
334                                    'lock_time', 'lock_reason', 'protection',])
335            self.print_by_ids(acls, 'ACLs', line_before=True)
336            labels = self._cleanup_labels(labels)
337            self.print_by_ids(labels, 'Labels', line_before=True)
338            self.print_dict(attributes, 'Host Attributes', line_before=True)
339
340
341class host_jobs(host):
342    """atest host jobs [--max-query] --mlist <file>|<hosts>"""
343    usage_action = 'jobs'
344
345    def __init__(self):
346        super(host_jobs, self).__init__()
347        self.parser.add_option('-q', '--max-query',
348                               help='Limits the number of results '
349                               '(20 by default)',
350                               type='int', default=20)
351
352
353    def parse(self):
354        """Consume the specific options"""
355        (options, leftover) = super(host_jobs, self).parse()
356        self.max_queries = options.max_query
357        return (options, leftover)
358
359
360    def execute(self):
361        """Execute 'atest host jobs'."""
362        results = []
363        real_hosts = []
364        for host in self.hosts:
365            if host.endswith('*'):
366                stats = self.execute_rpc('get_hosts',
367                                         hostname__startswith=host.rstrip('*'))
368                if len(stats) == 0:
369                    self.failure('No host matching %s' % host, item=host,
370                                 what_failed='Failed to stat')
371                [real_hosts.append(stat['hostname']) for stat in stats]
372            else:
373                real_hosts.append(host)
374
375        for host in real_hosts:
376            queue_entries = self.execute_rpc('get_host_queue_entries',
377                                             host__hostname=host,
378                                             query_limit=self.max_queries,
379                                             sort_by=['-job__id'])
380            jobs = []
381            for entry in queue_entries:
382                job = {'job_id': entry['job']['id'],
383                       'job_owner': entry['job']['owner'],
384                       'job_name': entry['job']['name'],
385                       'status': entry['status']}
386                jobs.append(job)
387            results.append((host, jobs))
388        return results
389
390
391    def output(self, results):
392        """Print output of 'atest host jobs'.
393
394        @param results: the results to be printed.
395        """
396        for host, jobs in results:
397            print '-'*5
398            print 'Hostname: %s' % host
399            self.print_table(jobs, keys_header=['job_id',
400                                                'job_owner',
401                                                'job_name',
402                                                'status'])
403
404class BaseHostModCreate(host):
405    """The base class for host_mod and host_create"""
406    # Matches one attribute=value pair
407    attribute_regex = r'(?P<attribute>\w+)=(?P<value>.+)?'
408
409    def __init__(self):
410        """Add the options shared between host mod and host create actions."""
411        self.messages = []
412        self.host_ids = {}
413        super(BaseHostModCreate, self).__init__()
414        self.parser.add_option('-l', '--lock',
415                               help='Lock hosts.',
416                               action='store_true')
417        self.parser.add_option('-r', '--lock_reason',
418                               help='Reason for locking hosts.',
419                               default='')
420        self.parser.add_option('-u', '--unlock',
421                               help='Unlock hosts.',
422                               action='store_true')
423
424        self.parser.add_option('-p', '--protection', type='choice',
425                               help=('Set the protection level on a host.  '
426                                     'Must be one of: %s. %s' %
427                                     (', '.join('"%s"' % p
428                                               for p in self.protections),
429                                      skylab_utils.MSG_INVALID_IN_SKYLAB)),
430                               choices=self.protections)
431        self._attributes = []
432        self.parser.add_option('--attribute', '-i',
433                               help=('Host attribute to add or change. Format '
434                                     'is <attribute>=<value>. Multiple '
435                                     'attributes can be set by passing the '
436                                     'argument multiple times. Attributes can '
437                                     'be unset by providing an empty value.'),
438                               action='append')
439        self.parser.add_option('-b', '--labels',
440                               help=('Comma separated list of labels. '
441                                     'When --skylab is provided, a label must '
442                                     'be in the format of label-key:label-value'
443                                     ' (e.g., board:lumpy).'))
444        self.parser.add_option('-B', '--blist',
445                               help='File listing the labels',
446                               type='string',
447                               metavar='LABEL_FLIST')
448        self.parser.add_option('-a', '--acls',
449                               help=('Comma separated list of ACLs. %s' %
450                                     skylab_utils.MSG_INVALID_IN_SKYLAB))
451        self.parser.add_option('-A', '--alist',
452                               help=('File listing the acls. %s' %
453                                     skylab_utils.MSG_INVALID_IN_SKYLAB),
454                               type='string',
455                               metavar='ACL_FLIST')
456        self.parser.add_option('-t', '--platform',
457                               help=('Sets the platform label. %s Please set '
458                                     'platform in labels (e.g., -b '
459                                     'platform:platform_name) with --skylab.' %
460                                     skylab_utils.MSG_INVALID_IN_SKYLAB))
461
462
463    def parse(self):
464        """Consume the options common to host create and host mod.
465        """
466        label_info = topic_common.item_parse_info(attribute_name='labels',
467                                                 inline_option='labels',
468                                                 filename_option='blist')
469        acl_info = topic_common.item_parse_info(attribute_name='acls',
470                                                inline_option='acls',
471                                                filename_option='alist')
472
473        (options, leftover) = super(BaseHostModCreate, self).parse([label_info,
474                                                              acl_info],
475                                                             req_items='hosts')
476
477        self._parse_lock_options(options)
478
479        self.label_map = None
480        if self.allow_skylab and self.skylab:
481            # TODO(nxia): drop these flags when all hosts are migrated to skylab
482            if (options.protection or options.acls or options.alist or
483                options.platform):
484                self.invalid_syntax(
485                        '--protection, --acls, --alist or --platform is not '
486                        'supported with --skylab.')
487
488            if self.labels:
489                self.label_map = device.convert_to_label_map(self.labels)
490
491        if options.protection:
492            self.data['protection'] = options.protection
493            self.messages.append('Protection set to "%s"' % options.protection)
494
495        self.attributes = {}
496        if options.attribute:
497            for pair in options.attribute:
498                m = re.match(self.attribute_regex, pair)
499                if not m:
500                    raise topic_common.CliError('Attribute must be in key=value '
501                                                'syntax.')
502                elif m.group('attribute') in self.attributes:
503                    raise topic_common.CliError(
504                            'Multiple values provided for attribute '
505                            '%s.' % m.group('attribute'))
506                self.attributes[m.group('attribute')] = m.group('value')
507
508        self.platform = options.platform
509        return (options, leftover)
510
511
512    def _set_acls(self, hosts, acls):
513        """Add hosts to acls (and remove from all other acls).
514
515        @param hosts: list of hostnames
516        @param acls: list of acl names
517        """
518        # Remove from all ACLs except 'Everyone' and ACLs in list
519        # Skip hosts that don't exist
520        for host in hosts:
521            if host not in self.host_ids:
522                continue
523            host_id = self.host_ids[host]
524            for a in self.execute_rpc('get_acl_groups', hosts=host_id):
525                if a['name'] not in self.acls and a['id'] != 1:
526                    self.execute_rpc('acl_group_remove_hosts', id=a['id'],
527                                     hosts=self.hosts)
528
529        # Add hosts to the ACLs
530        self.check_and_create_items('get_acl_groups', 'add_acl_group',
531                                    self.acls)
532        for a in acls:
533            self.execute_rpc('acl_group_add_hosts', id=a, hosts=hosts)
534
535
536    def _remove_labels(self, host, condition):
537        """Remove all labels from host that meet condition(label).
538
539        @param host: hostname
540        @param condition: callable that returns bool when given a label
541        """
542        if host in self.host_ids:
543            host_id = self.host_ids[host]
544            labels_to_remove = []
545            for l in self.execute_rpc('get_labels', host=host_id):
546                if condition(l):
547                    labels_to_remove.append(l['id'])
548            if labels_to_remove:
549                self.execute_rpc('host_remove_labels', id=host_id,
550                                 labels=labels_to_remove)
551
552
553    def _set_labels(self, host, labels):
554        """Apply labels to host (and remove all other labels).
555
556        @param host: hostname
557        @param labels: list of label names
558        """
559        condition = lambda l: l['name'] not in labels and not l['platform']
560        self._remove_labels(host, condition)
561        self.check_and_create_items('get_labels', 'add_label', labels)
562        self.execute_rpc('host_add_labels', id=host, labels=labels)
563
564
565    def _set_platform_label(self, host, platform_label):
566        """Apply the platform label to host (and remove existing).
567
568        @param host: hostname
569        @param platform_label: platform label's name
570        """
571        self._remove_labels(host, lambda l: l['platform'])
572        self.check_and_create_items('get_labels', 'add_label', [platform_label],
573                                    platform=True)
574        self.execute_rpc('host_add_labels', id=host, labels=[platform_label])
575
576
577    def _set_attributes(self, host, attributes):
578        """Set attributes on host.
579
580        @param host: hostname
581        @param attributes: attribute dictionary
582        """
583        for attr, value in self.attributes.iteritems():
584            self.execute_rpc('set_host_attribute', attribute=attr,
585                             value=value, hostname=host)
586
587
588class host_mod(BaseHostModCreate):
589    """atest host mod [--lock|--unlock --force_modify_locking
590    --platform <arch>
591    --labels <labels>|--blist <label_file>
592    --acls <acls>|--alist <acl_file>
593    --protection <protection_type>
594    --attributes <attr>=<value>;<attr>=<value>
595    --mlist <mach_file>] <hosts>"""
596    usage_action = 'mod'
597
598    def __init__(self):
599        """Add the options specific to the mod action"""
600        super(host_mod, self).__init__()
601        self.parser.add_option('--unlock-lock-id',
602                               help=('Unlock the lock with the lock-id. %s' %
603                                     skylab_utils.MSG_ONLY_VALID_IN_SKYLAB),
604                               default=None)
605        self.parser.add_option('-f', '--force_modify_locking',
606                               help='Forcefully lock\unlock a host',
607                               action='store_true')
608        self.parser.add_option('--remove_acls',
609                               help=('Remove all active acls. %s' %
610                                     skylab_utils.MSG_INVALID_IN_SKYLAB),
611                               action='store_true')
612        self.parser.add_option('--remove_labels',
613                               help='Remove all labels.',
614                               action='store_true')
615
616        self.add_skylab_options()
617        self.parser.add_option('--new-env',
618                               dest='new_env',
619                               choices=['staging', 'prod'],
620                               help=('The new environment ("staging" or '
621                                     '"prod") of the hosts. %s' %
622                                     skylab_utils.MSG_ONLY_VALID_IN_SKYLAB),
623                               default=None)
624
625
626    def _parse_unlock_options(self, options):
627        """Parse unlock related options."""
628        if self.skylab and options.unlock and options.unlock_lock_id is None:
629            self.invalid_syntax('Must provide --unlock-lock-id with "--skylab '
630                                '--unlock".')
631
632        if (not (self.skylab and options.unlock) and
633            options.unlock_lock_id is not None):
634            self.invalid_syntax('--unlock-lock-id is only valid with '
635                                '"--skylab --unlock".')
636
637        self.unlock_lock_id = options.unlock_lock_id
638
639
640    def parse(self):
641        """Consume the specific options"""
642        (options, leftover) = super(host_mod, self).parse()
643
644        self._parse_unlock_options(options)
645
646        if options.force_modify_locking:
647             self.data['force_modify_locking'] = True
648
649        if self.skylab and options.remove_acls:
650            # TODO(nxia): drop the flag when all hosts are migrated to skylab
651            self.invalid_syntax('--remove_acls is not supported with --skylab.')
652
653        self.remove_acls = options.remove_acls
654        self.remove_labels = options.remove_labels
655        self.new_env = options.new_env
656
657        return (options, leftover)
658
659
660    def execute_skylab(self):
661        """Execute atest host mod with --skylab.
662
663        @return A list of hostnames which have been successfully modified.
664        """
665        inventory_repo = skylab_utils.InventoryRepo(self.inventory_repo_dir)
666        inventory_repo.initialize()
667        data_dir = inventory_repo.get_data_dir()
668        lab = text_manager.load_lab(data_dir)
669
670        locked_by = None
671        if self.lock:
672            locked_by = inventory_repo.git_repo.config('user.email')
673
674        successes = []
675        for hostname in self.hosts:
676            try:
677                device.modify(
678                        lab,
679                        'duts',
680                        hostname,
681                        self.environment,
682                        lock=self.lock,
683                        locked_by=locked_by,
684                        lock_reason = self.lock_reason,
685                        unlock=self.unlock,
686                        unlock_lock_id=self.unlock_lock_id,
687                        attributes=self.attributes,
688                        remove_labels=self.remove_labels,
689                        label_map=self.label_map,
690                        new_env=self.new_env)
691                successes.append(hostname)
692            except device.SkylabDeviceActionError as e:
693                print('Cannot modify host %s: %s' % (hostname, e))
694
695        if successes:
696            text_manager.dump_lab(data_dir, lab)
697
698            status = inventory_repo.git_repo.status()
699            if not status:
700                print('Nothing is changed for hosts %s.' % successes)
701                return []
702
703            message = skylab_utils.construct_commit_message(
704                    'Modify %d hosts.\n\n%s' % (len(successes), successes))
705            self.change_number = inventory_repo.upload_change(
706                    message, draft=self.draft, dryrun=self.dryrun,
707                    submit=self.submit)
708
709        return successes
710
711
712    def execute(self):
713        """Execute 'atest host mod'."""
714        if self.skylab:
715            return self.execute_skylab()
716
717        successes = []
718        for host in self.execute_rpc('get_hosts', hostname__in=self.hosts):
719            self.host_ids[host['hostname']] = host['id']
720        for host in self.hosts:
721            if host not in self.host_ids:
722                self.failure('Cannot modify non-existant host %s.' % host)
723                continue
724            host_id = self.host_ids[host]
725
726            try:
727                if self.data:
728                    self.execute_rpc('modify_host', item=host,
729                                     id=host, **self.data)
730
731                if self.attributes:
732                    self._set_attributes(host, self.attributes)
733
734                if self.labels or self.remove_labels:
735                    self._set_labels(host, self.labels)
736
737                if self.platform:
738                    self._set_platform_label(host, self.platform)
739
740                # TODO: Make the AFE return True or False,
741                # especially for lock
742                successes.append(host)
743            except topic_common.CliError, full_error:
744                # Already logged by execute_rpc()
745                pass
746
747        if self.acls or self.remove_acls:
748            self._set_acls(self.hosts, self.acls)
749
750        return successes
751
752
753    def output(self, hosts):
754        """Print output of 'atest host mod'.
755
756        @param hosts: the host list to be printed.
757        """
758        for msg in self.messages:
759            self.print_wrapped(msg, hosts)
760
761        if hosts and self.skylab:
762            print('Modified hosts: %s.' % ', '.join(hosts))
763            if self.skylab and not self.dryrun and not self.submit:
764                print(skylab_utils.get_cl_message(self.change_number))
765
766
767class HostInfo(object):
768    """Store host information so we don't have to keep looking it up."""
769    def __init__(self, hostname, platform, labels):
770        self.hostname = hostname
771        self.platform = platform
772        self.labels = labels
773
774
775class host_create(BaseHostModCreate):
776    """atest host create [--lock|--unlock --platform <arch>
777    --labels <labels>|--blist <label_file>
778    --acls <acls>|--alist <acl_file>
779    --protection <protection_type>
780    --attributes <attr>=<value>;<attr>=<value>
781    --mlist <mach_file>] <hosts>"""
782    usage_action = 'create'
783
784    def parse(self):
785        """Option logic specific to create action.
786        """
787        (options, leftovers) = super(host_create, self).parse()
788        self.locked = options.lock
789        if 'serials' in self.attributes:
790            if len(self.hosts) > 1:
791                raise topic_common.CliError('Can not specify serials with '
792                                            'multiple hosts.')
793
794
795    @classmethod
796    def construct_without_parse(
797            cls, web_server, hosts, platform=None,
798            locked=False, lock_reason='', labels=[], acls=[],
799            protection=host_protections.Protection.NO_PROTECTION):
800        """Construct a host_create object and fill in data from args.
801
802        Do not need to call parse after the construction.
803
804        Return an object of site_host_create ready to execute.
805
806        @param web_server: A string specifies the autotest webserver url.
807            It is needed to setup comm to make rpc.
808        @param hosts: A list of hostnames as strings.
809        @param platform: A string or None.
810        @param locked: A boolean.
811        @param lock_reason: A string.
812        @param labels: A list of labels as strings.
813        @param acls: A list of acls as strings.
814        @param protection: An enum defined in host_protections.
815        """
816        obj = cls()
817        obj.web_server = web_server
818        try:
819            # Setup stuff needed for afe comm.
820            obj.afe = rpc.afe_comm(web_server)
821        except rpc.AuthError, s:
822            obj.failure(str(s), fatal=True)
823        obj.hosts = hosts
824        obj.platform = platform
825        obj.locked = locked
826        if locked and lock_reason.strip():
827            obj.data['lock_reason'] = lock_reason.strip()
828        obj.labels = labels
829        obj.acls = acls
830        if protection:
831            obj.data['protection'] = protection
832        obj.attributes = {}
833        return obj
834
835
836    def _detect_host_info(self, host):
837        """Detect platform and labels from the host.
838
839        @param host: hostname
840
841        @return: HostInfo object
842        """
843        # Mock an afe_host object so that the host is constructed as if the
844        # data was already in afe
845        data = {'attributes': self.attributes, 'labels': self.labels}
846        afe_host = frontend.Host(None, data)
847        store = host_info.InMemoryHostInfoStore(
848                host_info.HostInfo(labels=self.labels,
849                                   attributes=self.attributes))
850        machine = {
851                'hostname': host,
852                'afe_host': afe_host,
853                'host_info_store': store
854        }
855        try:
856            if bin_utils.ping(host, tries=1, deadline=1) == 0:
857                serials = self.attributes.get('serials', '').split(',')
858                adb_serial = self.attributes.get('serials')
859                host_dut = hosts.create_host(machine,
860                                             adb_serial=adb_serial)
861
862                info = HostInfo(host, host_dut.get_platform(),
863                                host_dut.get_labels())
864                # Clean host to make sure nothing left after calling it,
865                # e.g. tunnels.
866                if hasattr(host_dut, 'close'):
867                    host_dut.close()
868            else:
869                # Can't ping the host, use default information.
870                info = HostInfo(host, None, [])
871        except (socket.gaierror, error.AutoservRunError,
872                error.AutoservSSHTimeout):
873            # We may be adding a host that does not exist yet or we can't
874            # reach due to hostname/address issues or if the host is down.
875            info = HostInfo(host, None, [])
876        return info
877
878
879    def _execute_add_one_host(self, host):
880        # Always add the hosts as locked to avoid the host
881        # being picked up by the scheduler before it's ACL'ed.
882        self.data['locked'] = True
883        if not self.locked:
884            self.data['lock_reason'] = 'Forced lock on device creation'
885        self.execute_rpc('add_host', hostname=host, status="Ready", **self.data)
886
887        # If there are labels avaliable for host, use them.
888        info = self._detect_host_info(host)
889        labels = set(self.labels)
890        if info.labels:
891            labels.update(info.labels)
892
893        if labels:
894            self._set_labels(host, list(labels))
895
896        # Now add the platform label.
897        # If a platform was not provided and we were able to retrieve it
898        # from the host, use the retrieved platform.
899        platform = self.platform if self.platform else info.platform
900        if platform:
901            self._set_platform_label(host, platform)
902
903        if self.attributes:
904            self._set_attributes(host, self.attributes)
905
906
907    def execute(self):
908        """Execute 'atest host create'."""
909        successful_hosts = []
910        for host in self.hosts:
911            try:
912                self._execute_add_one_host(host)
913                successful_hosts.append(host)
914            except topic_common.CliError:
915                pass
916
917        if successful_hosts:
918            self._set_acls(successful_hosts, self.acls)
919
920            if not self.locked:
921                for host in successful_hosts:
922                    self.execute_rpc('modify_host', id=host, locked=False,
923                                     lock_reason='')
924        return successful_hosts
925
926
927    def output(self, hosts):
928        """Print output of 'atest host create'.
929
930        @param hosts: the added host list to be printed.
931        """
932        self.print_wrapped('Added host', hosts)
933
934
935class host_delete(action_common.atest_delete, host):
936    """atest host delete [--mlist <mach_file>] <hosts>"""
937
938    def __init__(self):
939        super(host_delete, self).__init__()
940
941        self.add_skylab_options()
942
943
944    def execute_skylab(self):
945        """Execute 'atest host delete' with '--skylab'.
946
947        @return A list of hostnames which have been successfully deleted.
948        """
949        inventory_repo = skylab_utils.InventoryRepo(self.inventory_repo_dir)
950        inventory_repo.initialize()
951        data_dir = inventory_repo.get_data_dir()
952        lab = text_manager.load_lab(data_dir)
953
954        successes = []
955        for hostname in self.hosts:
956            try:
957                device.delete(
958                        lab,
959                        'duts',
960                        hostname,
961                        self.environment)
962                successes.append(hostname)
963            except device.SkylabDeviceActionError as e:
964                print('Cannot delete host %s: %s' % (hostname, e))
965
966        if successes:
967            text_manager.dump_lab(data_dir, lab)
968            message = skylab_utils.construct_commit_message(
969                    'Delete %d hosts.\n\n%s' % (len(successes), successes))
970            self.change_number = inventory_repo.upload_change(
971                    message, draft=self.draft, dryrun=self.dryrun,
972                    submit=self.submit)
973
974        return successes
975
976
977    def execute(self):
978        """Execute 'atest host delete'.
979
980        @return A list of hostnames which have been successfully deleted.
981        """
982        if self.skylab:
983            return self.execute_skylab()
984
985        return super(host_delete, self).execute()
986
987
988class InvalidHostnameError(Exception):
989    """Cannot perform actions on the host because of invalid hostname."""
990
991
992def _add_hostname_suffix(hostname, suffix):
993    """Add the suffix to the hostname."""
994    if hostname.endswith(suffix):
995        raise InvalidHostnameError(
996              'Cannot add "%s" as it already contains the suffix.' % suffix)
997
998    return hostname + suffix
999
1000
1001def _remove_hostname_suffix(hostname, suffix):
1002    """Remove the suffix from the hostname."""
1003    if not hostname.endswith(suffix):
1004        raise InvalidHostnameError(
1005                'Cannot remove "%s" as it doesn\'t contain the suffix.' %
1006                suffix)
1007
1008    return hostname[:len(hostname) - len(suffix)]
1009
1010
1011class host_rename(host):
1012    """Host rename is only for migrating hosts between skylab and AFE DB."""
1013
1014    usage_action = 'rename'
1015
1016    def __init__(self):
1017        """Add the options specific to the rename action."""
1018        super(host_rename, self).__init__()
1019
1020        self.parser.add_option('--for-migration',
1021                               help=('Rename hostnames for migration. Rename '
1022                                     'each "hostname" to "hostname%s". '
1023                                     'The original "hostname" must not contain '
1024                                     'suffix.' % MIGRATED_HOST_SUFFIX),
1025                               action='store_true',
1026                               default=False)
1027        self.parser.add_option('--for-rollback',
1028                               help=('Rename hostnames for migration rollback. '
1029                                     'Rename each "hostname%s" to its original '
1030                                     '"hostname".' % MIGRATED_HOST_SUFFIX),
1031                               action='store_true',
1032                               default=False)
1033        self.parser.add_option('--dryrun',
1034                               help='Execute the action as a dryrun.',
1035                               action='store_true',
1036                               default=False)
1037
1038
1039    def parse(self):
1040        """Consume the options common to host rename."""
1041        (options, leftovers) = super(host_rename, self).parse()
1042        self.for_migration = options.for_migration
1043        self.for_rollback = options.for_rollback
1044        self.dryrun = options.dryrun
1045        self.host_ids = {}
1046
1047        if not (self.for_migration ^ self.for_rollback):
1048            self.invalid_syntax('--for-migration and --for-rollback are '
1049                                'exclusive, and one of them must be enabled.')
1050
1051        if not self.hosts:
1052            self.invalid_syntax('Must provide hostname(s).')
1053
1054        if self.dryrun:
1055            print('This will be a dryrun and will not rename hostnames.')
1056
1057        return (options, leftovers)
1058
1059
1060    def execute(self):
1061        """Execute 'atest host rename'."""
1062        if not self.prompt_confirmation():
1063            return
1064
1065        successes = []
1066        for host in self.execute_rpc('get_hosts', hostname__in=self.hosts):
1067            self.host_ids[host['hostname']] = host['id']
1068        for host in self.hosts:
1069            if host not in self.host_ids:
1070                self.failure('Cannot rename non-existant host %s.' % host,
1071                              item=host, what_failed='Failed to rename')
1072                continue
1073            try:
1074                host_id = self.host_ids[host]
1075                if self.for_migration:
1076                    new_hostname = _add_hostname_suffix(
1077                            host, MIGRATED_HOST_SUFFIX)
1078                else:
1079                    #for_rollback
1080                    new_hostname = _remove_hostname_suffix(
1081                            host, MIGRATED_HOST_SUFFIX)
1082
1083                if not self.dryrun:
1084                    # TODO(crbug.com/850737): delete and abort HQE.
1085                    data = {'hostname': new_hostname}
1086                    self.execute_rpc('modify_host', item=host, id=host_id,
1087                                     **data)
1088                successes.append((host, new_hostname))
1089            except InvalidHostnameError as e:
1090                self.failure('Cannot rename host %s: %s' % (host, e), item=host,
1091                             what_failed='Failed to rename')
1092            except topic_common.CliError, full_error:
1093                # Already logged by execute_rpc()
1094                pass
1095
1096        return successes
1097
1098
1099    def output(self, results):
1100        """Print output of 'atest host rename'."""
1101        if results:
1102            print('Successfully renamed:')
1103            for old_hostname, new_hostname in results:
1104                print('%s to %s' % (old_hostname, new_hostname))
1105
1106
1107class host_migrate(action_common.atest_list, host):
1108    """'atest host migrate' to migrate or rollback hosts."""
1109
1110    usage_action = 'migrate'
1111
1112    def __init__(self):
1113        super(host_migrate, self).__init__()
1114
1115        self.parser.add_option('--migration',
1116                               dest='migration',
1117                               help='Migrate the hosts to skylab.',
1118                               action='store_true',
1119                               default=False)
1120        self.parser.add_option('--rollback',
1121                               dest='rollback',
1122                               help='Rollback the hosts migrated to skylab.',
1123                               action='store_true',
1124                               default=False)
1125        self.parser.add_option('--model',
1126                               help='Model of the hosts to migrate.',
1127                               dest='model',
1128                               default=None)
1129        self.parser.add_option('--board',
1130                               help='Board of the hosts to migrate.',
1131                               dest='board',
1132                               default=None)
1133        self.parser.add_option('--pool',
1134                               help=('Pool of the hosts to migrate. Must '
1135                                     'specify --model for the pool.'),
1136                               dest='pool',
1137                               default=None)
1138
1139        self.add_skylab_options(enforce_skylab=True)
1140
1141
1142    def parse(self):
1143        """Consume the specific options"""
1144        (options, leftover) = super(host_migrate, self).parse()
1145
1146        self.migration = options.migration
1147        self.rollback = options.rollback
1148        self.model = options.model
1149        self.pool = options.pool
1150        self.board = options.board
1151        self.host_ids = {}
1152
1153        if not (self.migration ^ self.rollback):
1154            self.invalid_syntax('--migration and --rollback are exclusive, '
1155                                'and one of them must be enabled.')
1156
1157        if self.pool is not None and (self.model is None and
1158                                      self.board is None):
1159            self.invalid_syntax('Must provide --model or --board with --pool.')
1160
1161        if not self.hosts and not (self.model or self.board):
1162            self.invalid_syntax('Must provide hosts or --model or --board.')
1163
1164        return (options, leftover)
1165
1166
1167    def _remove_invalid_hostnames(self, hostnames, log_failure=False):
1168        """Remove hostnames with MIGRATED_HOST_SUFFIX.
1169
1170        @param hostnames: A list of hostnames.
1171        @param log_failure: Bool indicating whether to log invalid hostsnames.
1172
1173        @return A list of valid hostnames.
1174        """
1175        invalid_hostnames = set()
1176        for hostname in hostnames:
1177            if hostname.endswith(MIGRATED_HOST_SUFFIX):
1178                if log_failure:
1179                    self.failure('Cannot migrate host with suffix "%s" %s.' %
1180                                 (MIGRATED_HOST_SUFFIX, hostname),
1181                                 item=hostname, what_failed='Failed to rename')
1182                invalid_hostnames.add(hostname)
1183
1184        hostnames = list(set(hostnames) - invalid_hostnames)
1185
1186        return hostnames
1187
1188
1189    def execute(self):
1190        """Execute 'atest host migrate'."""
1191        hostnames = self._remove_invalid_hostnames(self.hosts, log_failure=True)
1192
1193        filters = {}
1194        check_results = {}
1195        if hostnames:
1196            check_results['hostname__in'] = 'hostname'
1197            if self.migration:
1198                filters['hostname__in'] = hostnames
1199            else:
1200                # rollback
1201                hostnames_with_suffix = [
1202                        _add_hostname_suffix(h, MIGRATED_HOST_SUFFIX)
1203                        for h in hostnames]
1204                filters['hostname__in'] = hostnames_with_suffix
1205        else:
1206            # TODO(nxia): add exclude_filter {'hostname__endswith':
1207            # MIGRATED_HOST_SUFFIX} for --migration
1208            if self.rollback:
1209                filters['hostname__endswith'] = MIGRATED_HOST_SUFFIX
1210
1211        labels = []
1212        if self.model:
1213            labels.append('model:%s' % self.model)
1214        if self.pool:
1215            labels.append('pool:%s' % self.pool)
1216        if self.board:
1217            labels.append('board:%s' % self.board)
1218
1219        if labels:
1220            if len(labels) == 1:
1221                filters['labels__name__in'] = labels
1222                check_results['labels__name__in'] = None
1223            else:
1224                filters['multiple_labels'] = labels
1225                check_results['multiple_labels'] = None
1226
1227        results = super(host_migrate, self).execute(
1228                op='get_hosts', filters=filters, check_results=check_results)
1229        hostnames = [h['hostname'] for h in results]
1230
1231        if self.migration:
1232            hostnames = self._remove_invalid_hostnames(hostnames)
1233        else:
1234            # rollback
1235            hostnames = [_remove_hostname_suffix(h, MIGRATED_HOST_SUFFIX)
1236                         for h in hostnames]
1237
1238        return self.execute_skylab_migration(hostnames)
1239
1240
1241    def assign_duts_to_drone(self, infra, devices, environment):
1242        """Assign uids of the devices to a random skylab drone.
1243
1244        @param infra: An instance of lab_pb2.Infrastructure.
1245        @param devices: A list of device_pb2.Device to be assigned to the drone.
1246        @param environment: 'staging' or 'prod'.
1247        """
1248        skylab_drones = skylab_server.get_servers(
1249                infra, environment, role='skylab_drone', status='primary')
1250
1251        if len(skylab_drones) == 0:
1252            raise device.SkylabDeviceActionError(
1253                'No skylab drone is found in primary status and staging '
1254                'environment. Please confirm there is at least one valid skylab'
1255                ' drone added in skylab inventory.')
1256
1257        for device in devices:
1258            # Randomly distribute each device to a skylab_drone.
1259            skylab_drone = random.choice(skylab_drones)
1260            skylab_server.add_dut_uids(skylab_drone, [device])
1261
1262
1263    def remove_duts_from_drone(self, infra, devices):
1264        """Remove uids of the devices from their skylab drones.
1265
1266        @param infra: An instance of lab_pb2.Infrastructure.
1267        @devices: A list of device_pb2.Device to be remove from the drone.
1268        """
1269        skylab_drones = skylab_server.get_servers(
1270                infra, 'staging', role='skylab_drone', status='primary')
1271
1272        for skylab_drone in skylab_drones:
1273            skylab_server.remove_dut_uids(skylab_drone, devices)
1274
1275
1276    def execute_skylab_migration(self, hostnames):
1277        """Execute migration in skylab_inventory.
1278
1279        @param hostnames: A list of hostnames to migrate.
1280        @return If there're hosts to migrate, return a list of the hostnames and
1281                a message instructing actions after the migration; else return
1282                None.
1283        """
1284        if not hostnames:
1285            return
1286
1287        inventory_repo = skylab_utils.InventoryRepo(self.inventory_repo_dir)
1288        inventory_repo.initialize()
1289
1290        subdirs = ['skylab', 'prod', 'staging']
1291        data_dirs = skylab_data_dir, prod_data_dir, staging_data_dir = [
1292                inventory_repo.get_data_dir(data_subdir=d) for d in subdirs]
1293        skylab_lab, prod_lab, staging_lab = [
1294                text_manager.load_lab(d) for d in data_dirs]
1295        infra = text_manager.load_infrastructure(skylab_data_dir)
1296
1297        label_map = None
1298        labels = []
1299        if self.board:
1300            labels.append('board:%s' % self.board)
1301        if self.model:
1302            labels.append('model:%s' % self.model)
1303        if self.pool:
1304            labels.append('critical_pool:%s' % self.pool)
1305        if labels:
1306            label_map = device.convert_to_label_map(labels)
1307
1308        if self.migration:
1309            prod_devices = device.move_devices(
1310                    prod_lab, skylab_lab, 'duts', label_map=label_map,
1311                    hostnames=hostnames)
1312            staging_devices = device.move_devices(
1313                    staging_lab, skylab_lab, 'duts', label_map=label_map,
1314                    hostnames=hostnames)
1315
1316            all_devices = prod_devices + staging_devices
1317            # Hostnames in afe_hosts tabel.
1318            device_hostnames = [str(d.common.hostname) for d in all_devices]
1319            message = (
1320                'Migration: move %s hosts into skylab_inventory.\n\n'
1321                'Please run this command after the CL is submitted:\n'
1322                'atest host rename --for-migration %s' %
1323                (len(all_devices), ' '.join(device_hostnames)))
1324
1325            self.assign_duts_to_drone(infra, prod_devices, 'prod')
1326            self.assign_duts_to_drone(infra, staging_devices, 'staging')
1327        else:
1328            # rollback
1329            prod_devices = device.move_devices(
1330                    skylab_lab, prod_lab, 'duts', environment='prod',
1331                    label_map=label_map, hostnames=hostnames)
1332            staging_devices = device.move_devices(
1333                    skylab_lab, staging_lab, 'duts', environment='staging',
1334                    label_map=label_map, hostnames=hostnames)
1335
1336            all_devices = prod_devices + staging_devices
1337            # Hostnames in afe_hosts tabel.
1338            device_hostnames = [_add_hostname_suffix(str(d.common.hostname),
1339                                                     MIGRATED_HOST_SUFFIX)
1340                                for d in all_devices]
1341            message = (
1342                'Rollback: remove %s hosts from skylab_inventory.\n\n'
1343                'Please run this command after the CL is submitted:\n'
1344                'atest host rename --for-rollback %s' %
1345                (len(all_devices), ' '.join(device_hostnames)))
1346
1347            self.remove_duts_from_drone(infra, all_devices)
1348
1349        if all_devices:
1350            text_manager.dump_infrastructure(skylab_data_dir, infra)
1351
1352            if prod_devices:
1353                text_manager.dump_lab(prod_data_dir, prod_lab)
1354
1355            if staging_devices:
1356                text_manager.dump_lab(staging_data_dir, staging_lab)
1357
1358            text_manager.dump_lab(skylab_data_dir, skylab_lab)
1359
1360            self.change_number = inventory_repo.upload_change(
1361                    message, draft=self.draft, dryrun=self.dryrun,
1362                    submit=self.submit)
1363
1364            return all_devices, message
1365
1366
1367    def output(self, result):
1368        """Print output of 'atest host list'.
1369
1370        @param result: the result to be printed.
1371        """
1372        if result:
1373            devices, message = result
1374
1375            if devices:
1376                hostnames = [h.common.hostname for h in devices]
1377                if self.migration:
1378                    print('Migrating hosts: %s' % ','.join(hostnames))
1379                else:
1380                    # rollback
1381                    print('Rolling back hosts: %s' % ','.join(hostnames))
1382
1383                if not self.dryrun:
1384                    if not self.submit:
1385                        print(skylab_utils.get_cl_message(self.change_number))
1386                    else:
1387                        # Print the instruction command for renaming hosts.
1388                        print('%s' % message)
1389        else:
1390            print('No hosts were migrated.')
1391