1#!/usr/bin/python2
2# Copyright 2018 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Command for viewing and changing software version assignments.
7
8Usage:
9    stable_version [ -w SERVER ] [ -n ] [ -t TYPE ]
10    stable_version [ -w SERVER ] [ -n ] [ -t TYPE ] BOARD/MODEL
11    stable_version [ -w SERVER ] [ -n ] -t TYPE -d BOARD/MODEL
12    stable_version [ -w SERVER ] [ -n ] -t TYPE BOARD/MODEL VERSION
13
14Available options:
15-w SERVER | --web SERVER
16    Used to specify an alternative server for the AFE RPC interface.
17
18-n | --dry-run
19    When specified, the command reports what would be done, but makes no
20    changes.
21
22-t TYPE | --type TYPE
23    Specifies the type of version mapping to use.  This option is
24    required for operations to change or delete mappings.  When listing
25    mappings, the option may be omitted, in which case all mapping types
26    are listed.
27
28-d | --delete
29    Delete the mapping for the given board or model argument.
30
31Command arguments:
32BOARD/MODEL
33    When specified, indicates the board or model to use as a key when
34    listing, changing, or deleting mappings.
35
36VERSION
37    When specified, indicates that the version name should be assigned
38    to the given board or model.
39
40With no arguments, the command will list all available mappings of all
41types.  The `--type` option will restrict the listing to only mappings of
42the given type.
43
44With only a board or model specified (and without the `--delete`
45option), will list all mappings for the given board or model.  The
46`--type` option will restrict the listing to only mappings of the given
47type.
48
49With the `--delete` option, will delete the mapping for the given board
50or model.  The `--type` option is required in this case.
51
52With both a board or model and a version specified, will assign the
53version to the given board or model.  The `--type` option is required in
54this case.
55"""
56
57import argparse
58import os
59import sys
60
61import common
62from autotest_lib.server import frontend
63from autotest_lib.site_utils.stable_images import build_data
64
65
66class _CommandError(Exception):
67    """Exception to indicate an error in command processing."""
68
69
70class _VersionMapHandler(object):
71    """An internal class to wrap data for version map operations.
72
73    This is a simple class to gather in one place data associated
74    with higher-level command line operations.
75
76    @property _description  A string description used to describe the
77                            image type when printing command output.
78    @property _dry_run      Value of the `--dry-run` command line
79                            operation.
80    @property _afe          AFE RPC object.
81    @property _version_map  AFE version map object for the image type.
82    """
83
84    # Subclasses are required to redefine both of these to a string with
85    # an appropriate value.
86    TYPE = None
87    DESCRIPTION = None
88
89    def __init__(self, afe, dry_run):
90        self._afe = afe
91        self._dry_run = dry_run
92        self._version_map = afe.get_stable_version_map(self.TYPE)
93
94    @property
95    def _description(self):
96        return self.DESCRIPTION
97
98    def _format_key_data(self, key):
99        return '%-10s %-12s' % (self._description, key)
100
101    def _format_operation(self, opname, key):
102        return '%-9s %s' % (opname, self._format_key_data(key))
103
104    def get_mapping(self, key):
105        """Return the mapping for `key`.
106
107        @param key  Board or model key to use for look up.
108        """
109        return self._version_map.get_version(key)
110
111    def print_all_mappings(self):
112        """Print all mappings in `self._version_map`"""
113        print '%s version mappings:' % self._description
114        mappings = self._version_map.get_all_versions()
115        if not mappings:
116            return
117        key_list = mappings.keys()
118        key_width = max(12, len(max(key_list, key=len)))
119        format = '%%-%ds  %%s' % key_width
120        for k in sorted(key_list):
121            print format % (k, mappings[k])
122
123    def print_mapping(self, key):
124        """Print the mapping for `key`.
125
126        Prints a single mapping for the board/model specified by
127        `key`.  Print nothing if no mapping exists.
128
129        @param key  Board or model key to use for look up.
130        """
131        version = self.get_mapping(key)
132        if version is not None:
133            print '%s  %s' % (self._format_key_data(key), version)
134
135    def set_mapping(self, key, new_version):
136        """Change the mapping for `key`, and report the action.
137
138        The mapping for the board or model specifed by `key` is set
139        to `new_version`.  The setting is reported to the user as
140        added, changed, or unchanged based on the current mapping in
141        the AFE.
142
143        This operation honors `self._dry_run`.
144
145        @param key          Board or model key for assignment.
146        @param new_version  Version to be assigned to `key`.
147        """
148        old_version = self.get_mapping(key)
149        if old_version is None:
150            print '%s -> %s' % (
151                self._format_operation('Adding', key), new_version)
152        elif old_version != new_version:
153            print '%s -> %s to %s' % (
154                self._format_operation('Updating', key),
155                old_version, new_version)
156        else:
157            print '%s -> %s' % (
158                self._format_operation('Unchanged', key), old_version)
159        if not self._dry_run and old_version != new_version:
160            self._version_map.set_version(key, new_version)
161
162    def delete_mapping(self, key):
163        """Delete the mapping for `key`, and report the action.
164
165        The mapping for the board or model specifed by `key` is removed
166        from `self._version_map`.  The change is reported to the user.
167
168        Requests to delete non-existent keys are ignored.
169
170        This operation honors `self._dry_run`.
171
172        @param key  Board or model key to be deleted.
173        """
174        version = self.get_mapping(key)
175        if version is not None:
176            print '%s -> %s' % (
177                self._format_operation('Delete', key), version)
178            if not self._dry_run:
179                self._version_map.delete_version(key)
180        else:
181            print self._format_operation('Unmapped', key)
182
183
184class _FirmwareVersionMapHandler(_VersionMapHandler):
185    TYPE = frontend.AFE.FIRMWARE_IMAGE_TYPE
186    DESCRIPTION = 'Firmware'
187
188
189class _CrOSVersionMapHandler(_VersionMapHandler):
190    TYPE = frontend.AFE.CROS_IMAGE_TYPE
191    DESCRIPTION = 'Chrome OS'
192
193    def set_mapping(self, board, version):
194        """Assign the Chrome OS mapping for the given board.
195
196        This function assigns the given Chrome OS version to the given
197        board.  Additionally, for any model with firmware bundled in the
198        assigned build, that model will be assigned the firmware version
199        found for it in the build.
200
201        @param board    Chrome OS board to be assigned a new version.
202        @param version  New Chrome OS version to be assigned to the
203                        board.
204        """
205        new_version = build_data.get_omaha_upgrade(
206            build_data.get_omaha_version_map(), board, version)
207        if new_version != version:
208            print 'Force %s version from Omaha:  %-12s -> %s' % (
209                self._description, board, new_version)
210        super(_CrOSVersionMapHandler, self).set_mapping(board, new_version)
211        fw_versions = build_data.get_firmware_versions(board, new_version)
212        fw_handler = _FirmwareVersionMapHandler(self._afe, self._dry_run)
213        for model, fw_version in fw_versions.iteritems():
214            if fw_version is not None:
215                fw_handler.set_mapping(model, fw_version)
216
217    def delete_mapping(self, board):
218        """Delete the Chrome OS mapping for the given board.
219
220        This function handles deletes the Chrome OS version mapping for the
221        given board.  Additionally, any R/W firmware mapping that existed
222        because of the OS mapping will be deleted as well.
223
224        @param board    Chrome OS board to be deleted from the mapping.
225        """
226        version = self.get_mapping(board)
227        super(_CrOSVersionMapHandler, self).delete_mapping(board)
228        fw_versions = build_data.get_firmware_versions(board, version)
229        fw_handler = _FirmwareVersionMapHandler(self._afe, self._dry_run)
230        for model in fw_versions.iterkeys():
231            fw_handler.delete_mapping(model)
232
233
234class _FAFTVersionMapHandler(_VersionMapHandler):
235    TYPE = frontend.AFE.FAFT_IMAGE_TYPE
236    DESCRIPTION = 'FAFT'
237
238
239_IMAGE_TYPE_CLASSES = [
240    _CrOSVersionMapHandler,
241    _FirmwareVersionMapHandler,
242    _FAFTVersionMapHandler,
243]
244_ALL_IMAGE_TYPES = [cls.TYPE for cls in _IMAGE_TYPE_CLASSES]
245_IMAGE_TYPE_HANDLERS = {cls.TYPE: cls for cls in _IMAGE_TYPE_CLASSES}
246
247
248def _create_version_map_handler(image_type, afe, dry_run):
249    return _IMAGE_TYPE_HANDLERS[image_type](afe, dry_run)
250
251
252def _requested_mapping_handlers(afe, image_type):
253    """Iterate through the image types for a listing operation.
254
255    When listing all mappings, or when listing by board, the listing can
256    be either for all available image types, or just for a single type
257    requested on the command line.
258
259    This function takes the value of the `-t` option, and yields a
260    `_VersionMapHandler` object for either the single requested type, or
261    for all of the types.
262
263    @param afe          AFE RPC interface object; created from SERVER.
264    @param image_type   Argument to the `-t` option.  A non-empty string
265                        indicates a single image type; value of `None`
266                        indicates all types.
267    """
268    if image_type:
269        yield _create_version_map_handler(image_type, afe, True)
270    else:
271        for cls in _IMAGE_TYPE_CLASSES:
272            yield cls(afe, True)
273
274
275def list_all_mappings(afe, image_type):
276    """List all mappings in the AFE.
277
278    This function handles the following syntax usage case:
279
280        stable_version [-w SERVER] [-t TYPE]
281
282    @param afe          AFE RPC interface object; created from SERVER.
283    @param image_type   Argument to the `-t` option.
284    """
285    need_newline = False
286    for handler in _requested_mapping_handlers(afe, image_type):
287        if need_newline:
288            print
289        handler.print_all_mappings()
290        need_newline = True
291
292
293def list_mapping_by_key(afe, image_type, key):
294    """List all mappings for the given board or model.
295
296    This function handles the following syntax usage case:
297
298        stable_version [-w SERVER] [-t TYPE] BOARD/MODEL
299
300    @param afe          AFE RPC interface object; created from SERVER.
301    @param image_type   Argument to the `-t` option.
302    @param key          Value of the BOARD/MODEL argument.
303    """
304    for handler in _requested_mapping_handlers(afe, image_type):
305        handler.print_mapping(key)
306
307
308def _validate_set_mapping(arguments):
309    """Validate syntactic requirements to assign a mapping.
310
311    The given arguments specified assigning version to be assigned to
312    a board or model; check the arguments for errors that can't be
313    discovered by `ArgumentParser`.  Errors are reported by raising
314    `_CommandError`.
315
316    @param arguments  `Namespace` object returned from argument parsing.
317    """
318    if not arguments.type:
319        raise _CommandError('The -t/--type option is required to assign a '
320                            'version')
321    if arguments.type == _FirmwareVersionMapHandler.TYPE:
322        msg = ('Cannot assign %s versions directly; '
323               'must assign the %s version instead.')
324        descriptions = (_FirmwareVersionMapHandler.DESCRIPTION,
325                        _CrOSVersionMapHandler.DESCRIPTION)
326        raise _CommandError(msg % descriptions)
327
328
329def set_mapping(afe, image_type, key, version, dry_run):
330    """Assign a version mapping to the given board or model.
331
332    This function handles the following syntax usage case:
333
334        stable_version [-w SERVER] [-n] -t TYPE BOARD/MODEL VERSION
335
336    @param afe          AFE RPC interface object; created from SERVER.
337    @param image_type   Argument to the `-t` option.
338    @param key          Value of the BOARD/MODEL argument.
339    @param key          Value of the VERSION argument.
340    @param dry_run      Whether the `-n` option was supplied.
341    """
342    if dry_run:
343        print 'Dry run; no mappings will be changed.'
344    handler = _create_version_map_handler(image_type, afe, dry_run)
345    handler.set_mapping(key, version)
346
347
348def _validate_delete_mapping(arguments):
349    """Validate syntactic requirements to delete a mapping.
350
351    The given arguments specified the `-d` / `--delete` option; check
352    the arguments for errors that can't be discovered by
353    `ArgumentParser`.  Errors are reported by raising `_CommandError`.
354
355    @param arguments  `Namespace` object returned from argument parsing.
356    """
357    if arguments.key is None:
358        raise _CommandError('Must specify BOARD_OR_MODEL argument '
359                            'with -d/--delete')
360    if arguments.version is not None:
361        raise _CommandError('Cannot specify VERSION argument with '
362                            '-d/--delete')
363    if not arguments.type:
364        raise _CommandError('-t/--type required with -d/--delete option')
365
366
367def delete_mapping(afe, image_type, key, dry_run):
368    """Delete the version mapping for the given board or model.
369
370    This function handles the following syntax usage case:
371
372        stable_version [-w SERVER] [-n] -t TYPE -d BOARD/MODEL
373
374    @param afe          AFE RPC interface object; created from SERVER.
375    @param image_type   Argument to the `-t` option.
376    @param key          Value of the BOARD/MODEL argument.
377    @param dry_run      Whether the `-n` option was supplied.
378    """
379    if dry_run:
380        print 'Dry run; no mappings will be deleted.'
381    handler = _create_version_map_handler(image_type, afe, dry_run)
382    handler.delete_mapping(key)
383
384
385def _parse_args(argv):
386    """Parse the given arguments according to the command syntax.
387
388    @param argv   Full argument vector, with argv[0] being the command
389                  name.
390    """
391    parser = argparse.ArgumentParser(
392        prog=os.path.basename(argv[0]),
393        description='Set and view software version assignments')
394    parser.add_argument('-w', '--web', default=None,
395                        metavar='SERVER',
396                        help='Specify the AFE to query.')
397    parser.add_argument('-n', '--dry-run', action='store_true',
398                        help='Report what would be done without making '
399                             'changes.')
400    parser.add_argument('-t', '--type', default=None,
401                        choices=_ALL_IMAGE_TYPES,
402                        help='Specify type of software image to be assigned.')
403    parser.add_argument('-d', '--delete', action='store_true',
404                        help='Delete the BOARD_OR_MODEL argument from the '
405                             'mappings.')
406    parser.add_argument('key', nargs='?', metavar='BOARD_OR_MODEL',
407                        help='Board, model, or other key for which to get or '
408                             'set a version')
409    parser.add_argument('version', nargs='?', metavar='VERSION',
410                        help='Version to be assigned')
411    return parser.parse_args(argv[1:])
412
413
414def _dispatch_command(afe, arguments):
415    if arguments.delete:
416        _validate_delete_mapping(arguments)
417        delete_mapping(afe, arguments.type, arguments.key,
418                       arguments.dry_run)
419    elif arguments.key is None:
420        list_all_mappings(afe, arguments.type)
421    elif arguments.version is None:
422        list_mapping_by_key(afe, arguments.type, arguments.key)
423    else:
424        _validate_set_mapping(arguments)
425        set_mapping(afe, arguments.type, arguments.key,
426                    arguments.version, arguments.dry_run)
427
428
429def main(argv):
430    """Standard main routine.
431
432    @param argv  Command line arguments including `sys.argv[0]`.
433    """
434    arguments = _parse_args(argv)
435    afe = frontend.AFE(server=arguments.web)
436    try:
437        _dispatch_command(afe, arguments)
438    except _CommandError as exc:
439        print >>sys.stderr, 'Error: %s' % str(exc)
440        sys.exit(1)
441
442
443if __name__ == '__main__':
444    try:
445        main(sys.argv)
446    except KeyboardInterrupt:
447        pass
448