# 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