1# Copyright 2015 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Argument processing for the DUT deployment tool.
6
7The argument processing is mostly a conventional client of
8`argparse`, except that if the command is invoked without required
9arguments, code here will start a line-oriented text dialog with the
10user to get the arguments.
11
12These are the arguments:
13  * (required) Board of the DUTs to be deployed.
14  * (required) Hostnames of the DUTs to be deployed.
15  * (optional) Version of the test image to be made the stable
16    repair image for the board to be deployed.  If omitted, the
17    existing setting is retained.
18
19The interactive dialog is invoked if the board and hostnames
20are omitted from the command line.
21
22"""
23
24import argparse
25import os
26import re
27import subprocess
28import sys
29import time
30
31
32# _BUILD_URI_FORMAT
33# A format template for a Google storage URI that designates
34# one build.  The template is to be filled in with a board
35# name and build version number.
36
37_BUILD_URI_FORMAT = 'gs://chromeos-image-archive/%s-release/%s'
38
39
40# _BUILD_PATTERNS
41# For user convenience, argument parsing allows various formats
42# for build version strings.  The function _normalize_build_name()
43# is used to convert the recognized syntaxes into the name as
44# it appears in Google storage.
45#
46# _BUILD_PATTERNS describe the recognized syntaxes for user-supplied
47# build versions, and information about how to convert them.  See the
48# normalize function for details.
49#
50# For user-supplied build versions, the following forms are supported:
51#   ####        - Indicates a canary; equivalent to ####.0.0.
52#   ####.#.#    - A full build version without the leading R##- prefix.
53#   R##-###.#.# - Canonical form of a build version.
54
55_BUILD_PATTERNS = [
56    (re.compile(r'^R\d+-\d+\.\d+\.\d+$'),   None),
57    (re.compile(r'^\d+\.\d+\.\d+$'),        'LATEST-%s'),
58    (re.compile(r'^\d+$'),                  'LATEST-%s.0.0'),
59]
60
61
62# _VALID_HOSTNAME_PATTERNS
63# A list of REs describing patterns that are acceptable as names
64# for DUTs in the test lab.  Names that don't match one of the
65# patterns will be rejected as invalid.
66
67_VALID_HOSTNAME_PATTERNS = [
68    re.compile(r'chromeos\d+-row\d+-rack\d+-host\d+')
69]
70
71
72def _build_path_exists(board, buildpath):
73    """Return whether a given build file exists in Google storage.
74
75    The `buildpath` refers to a specific file associated with
76    release builds for `board`.  The path may be one of the "LATEST"
77    files (e.g. "LATEST-7356.0.0"), or it could refer to a build
78    artifact (e.g. "R46-7356.0.0/image.zip").
79
80    The function constructs the full GS URI from the arguments, and
81    then tests for its existence with `gsutil ls`.
82
83    @param board        Board to be tested.
84    @param buildpath    Partial path of a file in Google storage.
85
86    @return Return a true value iff the designated file exists.
87    """
88    try:
89        gsutil_cmd = [
90                'gsutil', 'ls',
91                _BUILD_URI_FORMAT % (board, buildpath)
92        ]
93        status = subprocess.call(gsutil_cmd,
94                                 stdout=open('/dev/null', 'w'),
95                                 stderr=subprocess.STDOUT)
96        return status == 0
97    except:
98        return False
99
100
101def _normalize_build_name(board, build):
102    """Convert a user-supplied build version to canonical form.
103
104    Canonical form looks like  R##-####.#.#, e.g. R46-7356.0.0.
105    Acceptable user-supplied forms are describe under
106    _BUILD_PATTERNS, above.  The returned value will be the name of
107    a directory containing build artifacts from a release builder
108    for the board.
109
110    Walk through `_BUILD_PATTERNS`, trying to convert a user
111    supplied build version name into a directory name for valid
112    build artifacts.  Searching stops at the first pattern matched,
113    regardless of whether the designated build actually exists.
114
115    `_BUILD_PATTERNS` is a list of tuples.  The first element of the
116    tuple is an RE describing a valid user input.  The second
117    element of the tuple is a format pattern for a "LATEST" filename
118    in storage that can be used to obtain the full build version
119    associated with the user supplied version.  If the second element
120    is `None`, the user supplied build version is already in canonical
121    form.
122
123    @param board    Board to be tested.
124    @param build    User supplied version name.
125
126    @return Return the name of a directory in canonical form, or
127            `None` if the build doesn't exist.
128    """
129    for regex, fmt in _BUILD_PATTERNS:
130        if not regex.match(build):
131            continue
132        if fmt is not None:
133            try:
134                gsutil_cmd = [
135                    'gsutil', 'cat',
136                    _BUILD_URI_FORMAT % (board, fmt % build)
137                ]
138                return subprocess.check_output(
139                        gsutil_cmd, stderr=open('/dev/null', 'w'))
140            except:
141                return None
142        elif _build_path_exists(board, '%s/image.zip' % build):
143            return build
144        else:
145            return None
146    return None
147
148
149def _validate_board(board):
150    """Return whether a given board exists in Google storage.
151
152    For purposes of this function, a board exists if it has a
153    "LATEST-master" file in its release builder's directory.
154
155    N.B. For convenience, this function prints an error message
156    on stderr in certain failure cases.  This is currently useful
157    for argument processing, but isn't really ideal if the callers
158    were to get more complicated.
159
160    @param board    The board to be tested for existence.
161    @return Return a true value iff the board exists.
162    """
163    # In this case, the board doesn't exist, but we don't want
164    # an error message.
165    if board is None:
166        return False
167    # Check Google storage; report failures on stderr.
168    if _build_path_exists(board, 'LATEST-master'):
169        return True
170    else:
171        sys.stderr.write('Board %s doesn\'t exist.\n' % board)
172        return False
173
174
175def _validate_build(board, build):
176    """Return whether a given build exists in Google storage.
177
178    N.B. For convenience, this function prints an error message
179    on stderr in certain failure cases.  This is currently useful
180    for argument processing, but isn't really ideal if the callers
181    were to get more complicated.
182
183    @param board    The board to be tested for a build
184    @param build    The version of the build to be tested for.  This
185                    build may be in a user-specified (non-canonical)
186                    form.
187    @return If the given board+build exists, return its canonical
188            (normalized) version string.  If the build doesn't
189            exist, return a false value.
190    """
191    canonical_build = _normalize_build_name(board, build)
192    if not canonical_build:
193        sys.stderr.write(
194                'Build %s is not a valid build version for %s.\n' %
195                (build, board))
196    return canonical_build
197
198
199def _validate_hostname(hostname):
200    """Return whether a given hostname is valid for the test lab.
201
202    This is a sanity check meant to guarantee that host names follow
203    naming requirements for the test lab.
204
205    N.B. For convenience, this function prints an error message
206    on stderr in certain failure cases.  This is currently useful
207    for argument processing, but isn't really ideal if the callers
208    were to get more complicated.
209
210    @param hostname The host name to be checked.
211    @return Return a true value iff the hostname is valid.
212    """
213    for p in _VALID_HOSTNAME_PATTERNS:
214        if p.match(hostname):
215            return True
216    sys.stderr.write(
217            'Hostname %s doesn\'t match a valid location name.\n' %
218                hostname)
219    return False
220
221
222def _validate_arguments(arguments):
223    """Check command line arguments, and account for defaults.
224
225    Check that all command-line argument constraints are satisfied.
226    If errors are found, they are reported on `sys.stderr`.
227
228    If there are any fields with defined defaults that couldn't be
229    calculated when we constructed the argument parser, calculate
230    them now.
231
232    @param arguments  Parsed results from
233                      `ArgumentParser.parse_args()`.
234    @return Return `True` if there are no errors to report, or
235            `False` if there are.
236    """
237    if (not arguments.hostnames and
238            (arguments.board or arguments.build)):
239        sys.stderr.write(
240                'DUT hostnames are required with board or build.\n')
241        return False
242    if arguments.board is not None:
243        if not _validate_board(arguments.board):
244            return False
245        if (arguments.build is not None and
246                not _validate_build(arguments.board, arguments.build)):
247            return False
248    return True
249
250
251def _read_with_prompt(input, prompt):
252    """Print a prompt and then read a line of text.
253
254    @param input File-like object from which to read the line.
255    @param prompt String to print to stderr prior to reading.
256    @return Returns a string, stripped of whitespace.
257    """
258    full_prompt = '%s> ' % prompt
259    sys.stderr.write(full_prompt)
260    return input.readline().strip()
261
262
263def _read_board(input, default_board):
264    """Read a valid board name from user input.
265
266    Prompt the user to supply a board name, and read one line.  If
267    the line names a valid board, return the board name.  If the
268    line is blank and `default_board` is a non-empty string, returns
269    `default_board`.  Retry until a valid input is obtained.
270
271    `default_board` isn't checked; the caller is responsible for
272    ensuring its validity.
273
274    @param input          File-like object from which to read the
275                          board.
276    @param default_board  Value to return if the user enters a
277                          blank line.
278    @return Returns `default_board` or a validated board name.
279    """
280    if default_board:
281        board_prompt = 'board name [%s]' % default_board
282    else:
283        board_prompt = 'board name'
284    new_board = None
285    while not _validate_board(new_board):
286        new_board = _read_with_prompt(input, board_prompt).lower()
287        if new_board:
288            sys.stderr.write('Checking for valid board.\n')
289        elif default_board:
290            return default_board
291    return new_board
292
293
294def _read_build(input, board):
295    """Read a valid build version from user input.
296
297    Prompt the user to supply a build version, and read one line.
298    If the line names an existing version for the given board,
299    return the canonical build version.  If the line is blank,
300    return `None` (indicating the build shouldn't change).
301
302    @param input    File-like object from which to read the build.
303    @param board    Board for the build.
304    @return Returns canonical build version, or `None`.
305    """
306    build = False
307    prompt = 'build version (optional)'
308    while not build:
309        build = _read_with_prompt(input, prompt)
310        if not build:
311            return None
312        sys.stderr.write('Checking for valid build.\n')
313        build = _validate_build(board, build)
314    return build
315
316
317def _read_hostnames(input):
318    """Read a list of host names from user input.
319
320    Prompt the user to supply a list of host names.  Any number of
321    lines are allowed; input is terminated at the first blank line.
322    Any number of hosts names are allowed on one line.  Names are
323    separated by whitespace.
324
325    Only valid host names are accepted.  Invalid host names are
326    ignored, and a warning is printed.
327
328    @param input    File-like object from which to read the names.
329    @return Returns a list of validated host names.
330    """
331    hostnames = []
332    y_n = 'yes'
333    while not 'no'.startswith(y_n):
334        sys.stderr.write('enter hosts (blank line to end):\n')
335        while True:
336            new_hosts = input.readline().strip().split()
337            if not new_hosts:
338                break
339            for h in new_hosts:
340                if _validate_hostname(h):
341                    hostnames.append(h)
342        if not hostnames:
343            sys.stderr.write('Must provide at least one hostname.\n')
344            continue
345        prompt = 'More hosts? [y/N]'
346        y_n = _read_with_prompt(input, prompt).lower() or 'no'
347    return hostnames
348
349
350def _read_arguments(input, arguments):
351    """Dialog to read all needed arguments from the user.
352
353    The user is prompted in turn for a board, a build, and
354    hostnames.  Responses are stored in `arguments`.  The user is
355    given opportunity to accept or reject the responses before
356    continuing.
357
358    @param input      File-like object from which to read user
359                      responses.
360    @param arguments  Arguments object returned from
361                      `ArgumentParser.parse_args()`.  Results are
362                      stored here.
363    """
364    y_n = 'no'
365    while not 'yes'.startswith(y_n):
366        arguments.board = _read_board(input, arguments.board)
367        arguments.build = _read_build(input, arguments.board)
368        prompt = '%s build %s? [Y/n]' % (
369                arguments.board, arguments.build)
370        y_n = _read_with_prompt(input, prompt).lower() or 'yes'
371    arguments.hostnames = _read_hostnames(input)
372
373
374def _parse(argv):
375    """Parse the command line arguments.
376
377    Create an argument parser for this command's syntax, parse the
378    command line, and return the result of the ArgumentParser
379    parse_args() method.
380
381    @param argv Standard command line argument vector; argv[0] is
382                assumed to be the command name.
383    @return Result returned by ArgumentParser.parse_args().
384    """
385    parser = argparse.ArgumentParser(
386            prog=argv[0],
387            description='Install a test image on newly deployed DUTs')
388    # frontend.AFE(server=None) will use the default web server,
389    # so default for --web is `None`.
390    parser.add_argument('-w', '--web', metavar='SERVER', default=None,
391                        help='specify web server')
392    parser.add_argument('-d', '--dir',
393                        help='directory for logs')
394    parser.add_argument('-i', '--build',
395                        help='select stable test build version')
396    parser.add_argument('-n', '--noinstall', action='store_true',
397                        help='skip install (for script testing)')
398    parser.add_argument('-s', '--nostage', action='store_true',
399                        help='skip staging test image (for script testing)')
400    parser.add_argument('-t', '--nostable', action='store_true',
401                        help='skip changing stable test image '
402                             '(for script testing)')
403    parser.add_argument('board', nargs='?', metavar='BOARD',
404                        help='board for DUTs to be installed')
405    parser.add_argument('hostnames', nargs='*', metavar='HOSTNAME',
406                        help='host names of DUTs to be installed')
407    return parser.parse_args(argv[1:])
408
409
410def parse_command(argv, full_deploy):
411    """Get arguments for install from `argv` or the user.
412
413    Create an argument parser for this command's syntax, parse the
414    command line, and return the result of the ArgumentParser
415    parse_args() method.
416
417    If mandatory arguments are missing, execute a dialog with the
418    user to read the arguments from `sys.stdin`.  Fill in the
419    return value with the values read prior to returning.
420
421    @param argv         Standard command line argument vector;
422                        argv[0] is assumed to be the command name.
423    @param full_deploy  Whether this is for full deployment or
424                        repair.
425
426    @return Result, as returned by ArgumentParser.parse_args().
427    """
428    arguments = _parse(argv)
429    arguments.full_deploy = full_deploy
430    if arguments.board is None:
431        _read_arguments(sys.stdin, arguments)
432    elif not _validate_arguments(arguments):
433        return None
434    if not arguments.dir:
435        basename = 'deploy.%s.%s' % (
436                time.strftime('%Y-%m-%d.%H:%M:%S'),
437                arguments.board)
438        arguments.dir = os.path.join(os.environ['HOME'],
439                                     'Documents', basename)
440        os.makedirs(arguments.dir)
441    elif not os.path.isdir(arguments.dir):
442        os.mkdir(arguments.dir)
443    return arguments
444