# Copyright 2015 The Chromium OS Authors. All rights reserved. # Use of this source code is governed by a BSD-style license that can be # found in the LICENSE file. """Argument processing for the DUT deployment tool. The argument processing is mostly a conventional client of `argparse`, except that if the command is invoked without required arguments, code here will start a line-oriented text dialog with the user to get the arguments. These are the arguments: * (required) Board of the DUTs to be deployed. * (required) Hostnames of the DUTs to be deployed. * (optional) Version of the test image to be made the stable repair image for the board to be deployed. If omitted, the existing setting is retained. The interactive dialog is invoked if the board and hostnames are omitted from the command line. """ import argparse import collections import csv import datetime import os import re import subprocess import sys import dateutil.tz import common from autotest_lib.server.hosts import servo_host # _BUILD_URI_FORMAT # A format template for a Google storage URI that designates # one build. The template is to be filled in with a board # name and build version number. _BUILD_URI_FORMAT = 'gs://chromeos-image-archive/%s-release/%s' # _BUILD_PATTERNS # For user convenience, argument parsing allows various formats # for build version strings. The function _normalize_build_name() # is used to convert the recognized syntaxes into the name as # it appears in Google storage. # # _BUILD_PATTERNS describe the recognized syntaxes for user-supplied # build versions, and information about how to convert them. See the # normalize function for details. # # For user-supplied build versions, the following forms are supported: # #### - Indicates a canary; equivalent to ####.0.0. # ####.#.# - A full build version without the leading R##- prefix. # R##-###.#.# - Canonical form of a build version. _BUILD_PATTERNS = [ (re.compile(r'^R\d+-\d+\.\d+\.\d+$'), None), (re.compile(r'^\d+\.\d+\.\d+$'), 'LATEST-%s'), (re.compile(r'^\d+$'), 'LATEST-%s.0.0'), ] # _VALID_HOSTNAME_PATTERNS # A list of REs describing patterns that are acceptable as names # for DUTs in the test lab. Names that don't match one of the # patterns will be rejected as invalid. _VALID_HOSTNAME_PATTERNS = [ re.compile(r'chromeos\d+-row\d+-rack\d+-host\d+') ] # _EXPECTED_NUMBER_OF_HOST_INFO # The number of items per line when parsing the hostname_file csv file. _EXPECTED_NUMBER_OF_HOST_INFO = 8 # HostInfo # Namedtuple to store host info for processing when creating host in the afe. HostInfo = collections.namedtuple('HostInfo', ['hostname', 'host_attr_dict']) def _build_path_exists(board, buildpath): """Return whether a given build file exists in Google storage. The `buildpath` refers to a specific file associated with release builds for `board`. The path may be one of the "LATEST" files (e.g. "LATEST-7356.0.0"), or it could refer to a build artifact (e.g. "R46-7356.0.0/image.zip"). The function constructs the full GS URI from the arguments, and then tests for its existence with `gsutil ls`. @param board Board to be tested. @param buildpath Partial path of a file in Google storage. @return Return a true value iff the designated file exists. """ try: gsutil_cmd = [ 'gsutil', 'ls', _BUILD_URI_FORMAT % (board, buildpath) ] status = subprocess.call(gsutil_cmd, stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT) return status == 0 except: return False def _normalize_build_name(board, build): """Convert a user-supplied build version to canonical form. Canonical form looks like R##-####.#.#, e.g. R46-7356.0.0. Acceptable user-supplied forms are describe under _BUILD_PATTERNS, above. The returned value will be the name of a directory containing build artifacts from a release builder for the board. Walk through `_BUILD_PATTERNS`, trying to convert a user supplied build version name into a directory name for valid build artifacts. Searching stops at the first pattern matched, regardless of whether the designated build actually exists. `_BUILD_PATTERNS` is a list of tuples. The first element of the tuple is an RE describing a valid user input. The second element of the tuple is a format pattern for a "LATEST" filename in storage that can be used to obtain the full build version associated with the user supplied version. If the second element is `None`, the user supplied build version is already in canonical form. @param board Board to be tested. @param build User supplied version name. @return Return the name of a directory in canonical form, or `None` if the build doesn't exist. """ for regex, fmt in _BUILD_PATTERNS: if not regex.match(build): continue if fmt is not None: try: gsutil_cmd = [ 'gsutil', 'cat', _BUILD_URI_FORMAT % (board, fmt % build) ] return subprocess.check_output( gsutil_cmd, stderr=open('/dev/null', 'w')) except: return None elif _build_path_exists(board, '%s/image.zip' % build): return build else: return None return None def _validate_board(board): """Return whether a given board exists in Google storage. For purposes of this function, a board exists if it has a "LATEST-master" file in its release builder's directory. N.B. For convenience, this function prints an error message on stderr in certain failure cases. This is currently useful for argument processing, but isn't really ideal if the callers were to get more complicated. @param board The board to be tested for existence. @return Return a true value iff the board exists. """ # In this case, the board doesn't exist, but we don't want # an error message. if board is None: return False # Check Google storage; report failures on stderr. if _build_path_exists(board, 'LATEST-master'): return True else: sys.stderr.write('Board %s doesn\'t exist.\n' % board) return False def _validate_build(board, build): """Return whether a given build exists in Google storage. N.B. For convenience, this function prints an error message on stderr in certain failure cases. This is currently useful for argument processing, but isn't really ideal if the callers were to get more complicated. @param board The board to be tested for a build @param build The version of the build to be tested for. This build may be in a user-specified (non-canonical) form. @return If the given board+build exists, return its canonical (normalized) version string. If the build doesn't exist, return a false value. """ canonical_build = _normalize_build_name(board, build) if not canonical_build: sys.stderr.write( 'Build %s is not a valid build version for %s.\n' % (build, board)) return canonical_build def _validate_hostname(hostname): """Return whether a given hostname is valid for the test lab. This is a sanity check meant to guarantee that host names follow naming requirements for the test lab. N.B. For convenience, this function prints an error message on stderr in certain failure cases. This is currently useful for argument processing, but isn't really ideal if the callers were to get more complicated. @param hostname The host name to be checked. @return Return a true value iff the hostname is valid. """ for p in _VALID_HOSTNAME_PATTERNS: if p.match(hostname): return True sys.stderr.write( 'Hostname %s doesn\'t match a valid location name.\n' % hostname) return False def _is_hostname_file_valid(hostname_file): """Check that the hostname file is valid. The hostname file is deemed valid if: - the file exists. - the file is non-empty. @param hostname_file Filename of the hostname file to check. @return `True` if the hostname file is valid, False otherse. """ return os.path.exists(hostname_file) and os.path.getsize(hostname_file) > 0 def _validate_arguments(arguments): """Check command line arguments, and account for defaults. Check that all command-line argument constraints are satisfied. If errors are found, they are reported on `sys.stderr`. If there are any fields with defined defaults that couldn't be calculated when we constructed the argument parser, calculate them now. @param arguments Parsed results from `ArgumentParser.parse_args()`. @return Return `True` if there are no errors to report, or `False` if there are. """ # If both hostnames and hostname_file are specified, complain about that. if arguments.hostnames and arguments.hostname_file: sys.stderr.write( 'DUT hostnames and hostname file both specified, only ' 'specify one or the other.\n') return False if (arguments.hostname_file and not _is_hostname_file_valid(arguments.hostname_file)): sys.stderr.write( 'Specified hostname file must exist and be non-empty.\n') return False if (not arguments.hostnames and not arguments.hostname_file and (arguments.board or arguments.build)): sys.stderr.write( 'DUT hostnames are required with board or build.\n') return False if arguments.board is not None: if not _validate_board(arguments.board): return False if (arguments.build is not None and not _validate_build(arguments.board, arguments.build)): return False return True def _read_with_prompt(input, prompt): """Print a prompt and then read a line of text. @param input File-like object from which to read the line. @param prompt String to print to stderr prior to reading. @return Returns a string, stripped of whitespace. """ full_prompt = '%s> ' % prompt sys.stderr.write(full_prompt) return input.readline().strip() def _read_board(input, default_board): """Read a valid board name from user input. Prompt the user to supply a board name, and read one line. If the line names a valid board, return the board name. If the line is blank and `default_board` is a non-empty string, returns `default_board`. Retry until a valid input is obtained. `default_board` isn't checked; the caller is responsible for ensuring its validity. @param input File-like object from which to read the board. @param default_board Value to return if the user enters a blank line. @return Returns `default_board` or a validated board name. """ if default_board: board_prompt = 'board name [%s]' % default_board else: board_prompt = 'board name' new_board = None while not _validate_board(new_board): new_board = _read_with_prompt(input, board_prompt).lower() if new_board: sys.stderr.write('Checking for valid board.\n') elif default_board: return default_board return new_board def _read_build(input, board): """Read a valid build version from user input. Prompt the user to supply a build version, and read one line. If the line names an existing version for the given board, return the canonical build version. If the line is blank, return `None` (indicating the build shouldn't change). @param input File-like object from which to read the build. @param board Board for the build. @return Returns canonical build version, or `None`. """ build = False prompt = 'build version (optional)' while not build: build = _read_with_prompt(input, prompt) if not build: return None sys.stderr.write('Checking for valid build.\n') build = _validate_build(board, build) return build def _read_hostnames(input): """Read a list of host names from user input. Prompt the user to supply a list of host names. Any number of lines are allowed; input is terminated at the first blank line. Any number of hosts names are allowed on one line. Names are separated by whitespace. Only valid host names are accepted. Invalid host names are ignored, and a warning is printed. @param input File-like object from which to read the names. @return Returns a list of validated host names. """ hostnames = [] y_n = 'yes' while not 'no'.startswith(y_n): sys.stderr.write('enter hosts (blank line to end):\n') while True: new_hosts = input.readline().strip().split() if not new_hosts: break for h in new_hosts: if _validate_hostname(h): hostnames.append(h) if not hostnames: sys.stderr.write('Must provide at least one hostname.\n') continue prompt = 'More hosts? [y/N]' y_n = _read_with_prompt(input, prompt).lower() or 'no' return hostnames def _read_arguments(input, arguments): """Dialog to read all needed arguments from the user. The user is prompted in turn for a board, a build, and hostnames. Responses are stored in `arguments`. The user is given opportunity to accept or reject the responses before continuing. @param input File-like object from which to read user responses. @param arguments Namespace object returned from `ArgumentParser.parse_args()`. Results are stored here. """ y_n = 'no' while not 'yes'.startswith(y_n): arguments.board = _read_board(input, arguments.board) arguments.build = _read_build(input, arguments.board) prompt = '%s build %s? [Y/n]' % ( arguments.board, arguments.build) y_n = _read_with_prompt(input, prompt).lower() or 'yes' arguments.hostnames = _read_hostnames(input) def get_default_logdir_name(arguments): """Get default log directory name. @param arguments Namespace object returned from argument parsing. @return A filename as a string. """ return '{time}-{board}'.format( time=arguments.start_time.isoformat(), board=arguments.board) class _ArgumentParser(argparse.ArgumentParser): """ArgumentParser extended with boolean option pairs.""" # Arguments required when adding an option pair. _REQUIRED_PAIR_ARGS = {'dest', 'default'} def add_argument_pair(self, yes_flags, no_flags, **kwargs): """Add a pair of argument flags for a boolean option. @param yes_flags Iterable of flags to turn option on. May also be a single string. @param no_flags Iterable of flags to turn option off. May also be a single string. @param *kwargs Other arguments to pass to add_argument() """ missing_args = self._REQUIRED_PAIR_ARGS - set(kwargs) if missing_args: raise ValueError("Argument pair must have explicit %s" % (', '.join(missing_args),)) if isinstance(yes_flags, (str, unicode)): yes_flags = [yes_flags] if isinstance(no_flags, (str, unicode)): no_flags = [no_flags] self.add_argument(*yes_flags, action='store_true', **kwargs) self.add_argument(*no_flags, action='store_false', **kwargs) def _make_common_parser(command_name): """Create argument parser for common arguments. @param command_name The command name. @return ArgumentParser instance. """ parser = _ArgumentParser( prog=command_name, description='Install a test image on newly deployed DUTs') # frontend.AFE(server=None) will use the default web server, # so default for --web is `None`. parser.add_argument('-w', '--web', metavar='SERVER', default=None, help='specify web server') parser.add_argument('-d', '--dir', dest='logdir', help='directory for logs') parser.add_argument('-i', '--build', help='select stable test build version') parser.add_argument('-n', '--noinstall', action='store_true', help='skip install (for script testing)') parser.add_argument('-s', '--nostage', action='store_true', help='skip staging test image (for script testing)') parser.add_argument('-t', '--nostable', action='store_true', help='skip changing stable test image ' '(for script testing)') parser.add_argument('-f', '--hostname_file', help='CSV file that contains a list of hostnames and ' 'their details to install with.') parser.add_argument('board', nargs='?', metavar='BOARD', help='board for DUTs to be installed') parser.add_argument('hostnames', nargs='*', metavar='HOSTNAME', help='host names of DUTs to be installed') return parser def _add_upload_argument_pair(parser, default): """Add option pair for uploading logs. @param parser _ArgumentParser instance. @param default Default option value. """ parser.add_argument_pair('--upload', '--noupload', dest='upload', default=default, help='upload logs to GS bucket',) def _parse_hostname_file_line(hostname_file_row): """ Parse a line from the hostname_file and return a dict of the info. @param hostname_file_row: List of strings from each line in the hostname file. @returns a NamedTuple of (hostname, host_attr_dict). host_attr_dict is a dict of host attributes for the host. """ if len(hostname_file_row) != _EXPECTED_NUMBER_OF_HOST_INFO: raise Exception('hostname_file line has unexpected number of items ' '%d (expect %d): %s' % (len(hostname_file_row), _EXPECTED_NUMBER_OF_HOST_INFO, hostname_file_row)) # The file will have the info in the following order: # 0: board # 1: dut hostname # 2: dut/v4 mac address # 3: dut ip # 4: labstation hostname # 5: servo serial # 6: servo mac address # 7: servo ip return HostInfo( hostname=hostname_file_row[1], host_attr_dict={servo_host.SERVO_HOST_ATTR: hostname_file_row[4], servo_host.SERVO_SERIAL_ATTR: hostname_file_row[5]}) def parse_hostname_file(hostname_file): """ Parse the hostname_file and return a list of dicts for each line. @param hostname_file: CSV file that contains all the goodies. @returns a list of dicts where each line is broken down into a dict. """ host_info_list = [] # First line will be the header, no need to parse that. first_line_skipped = False with open(hostname_file) as f: hostname_file_reader = csv.reader(f) for row in hostname_file_reader: if not first_line_skipped: first_line_skipped = True continue host_info_list.append(_parse_hostname_file_line(row)) return host_info_list def parse_command(argv, full_deploy): """Get arguments for install from `argv` or the user. Create an argument parser for this command's syntax, parse the command line, and return the result of the ArgumentParser parse_args() method. If mandatory arguments are missing, execute a dialog with the user to read the arguments from `sys.stdin`. Fill in the return value with the values read prior to returning. @param argv Standard command line argument vector; argv[0] is assumed to be the command name. @param full_deploy Whether this is for full deployment or repair. @return Result, as returned by ArgumentParser.parse_args(). """ command_name = os.path.basename(argv[0]) parser = _make_common_parser(command_name) _add_upload_argument_pair(parser, default=full_deploy) arguments = parser.parse_args(argv[1:]) arguments.full_deploy = full_deploy if arguments.board is None: _read_arguments(sys.stdin, arguments) elif not _validate_arguments(arguments): return None arguments.start_time = datetime.datetime.now(dateutil.tz.tzlocal()) if not arguments.logdir: basename = get_default_logdir_name(arguments) arguments.logdir = os.path.join(os.environ['HOME'], 'Documents', basename) os.makedirs(arguments.logdir) elif not os.path.isdir(arguments.logdir): os.mkdir(arguments.logdir) if arguments.hostname_file: # Populate arguments.hostnames with the hostnames from the file. hostname_file_info_list = parse_hostname_file(arguments.hostname_file) arguments.hostnames = [host_info.hostname for host_info in hostname_file_info_list] arguments.host_info_list = hostname_file_info_list else: arguments.host_info_list = [] return arguments