Home | History | Annotate | Download | only in prepare 
      1 #!/usr/bin/python -u
      2 # Copyright 2019 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 """Tool to (re)prepare a DUT for lab deployment."""
      7 
      8 from __future__ import absolute_import
      9 from __future__ import division
     10 from __future__ import print_function
     11 
     12 import argparse
     13 import errno
     14 import logging
     15 import logging.config
     16 import os
     17 
     18 import common
     19 from autotest_lib.server import afe_utils
     20 from autotest_lib.server.hosts import file_store
     21 from autotest_lib.site_utils.deployment.prepare import dut as preparedut
     22 
     23 
     24 class DutPreparationError(Exception):
     25   """Generic error raised during DUT preparation."""
     26 
     27 
     28 def main():
     29   """Tool to (re)prepare a DUT for lab deployment."""
     30   opts = _parse_args()
     31   _configure_logging('prepare_dut', os.path.join(opts.results_dir, _LOG_FILE))
     32 
     33   info = _read_store(opts.host_info_file)
     34   repair_image = _get_cros_repair_image_name(info.board)
     35   logging.info('Using repair image %s, obtained from AFE', repair_image)
     36   with _create_host(opts.hostname, info, opts.results_dir) as host:
     37     if opts.dry_run:
     38       logging.info('DRY RUN: Would have run actions %s', opts.actions)
     39       return
     40 
     41     if 'stage-usb' in opts.actions:
     42       preparedut.download_image_to_servo_usb(host, repair_image)
     43     if 'install-firmware' in opts.actions:
     44       preparedut.install_firmware(host, opts.force_firmware)
     45     if 'install-test-image' in opts.actions:
     46       preparedut.install_test_image(host)
     47 
     48 
     49 _LOG_FILE = 'prepare_dut.log'
     50 _DUT_LOGS_DIR = 'dut_logs'
     51 
     52 
     53 def _parse_args():
     54   parser = argparse.ArgumentParser(
     55       description='Prepare / validate DUT for lab deployment.')
     56 
     57   parser.add_argument(
     58       'actions',
     59       nargs='+',
     60       choices=['stage-usb', 'install-firmware', 'install-test-image'],
     61       help='DUT preparation actions to execute.',
     62   )
     63   parser.add_argument(
     64       '--dry-run',
     65       action='store_true',
     66       default=False,
     67       help='Run in dry-run mode. No changes will be made to the DUT.',
     68   )
     69   parser.add_argument(
     70       '--results-dir',
     71       required=True,
     72       help='Directory to drop logs and output artifacts in.',
     73   )
     74 
     75   parser.add_argument(
     76       '--hostname',
     77       required=True,
     78       help='Hostname of the DUT to prepare.',
     79   )
     80   parser.add_argument(
     81       '--host-info-file',
     82       required=True,
     83       help=('Full path to HostInfo file.'
     84             ' DUT inventory information is read from the HostInfo file.'),
     85   )
     86 
     87   parser.add_argument(
     88       '--force-firmware',
     89       action='store_true',
     90       help='Force firmware isntallation via chromeos-installfirmware.',
     91   )
     92 
     93   return parser.parse_args()
     94 
     95 
     96 def _configure_logging(name, tee_file):
     97     """Configure logging globally.
     98 
     99     @param name: Name to prepend to log messages.
    100                  This should be the name of the program.
    101     @param tee_file: File to tee logs to, in addition to stderr.
    102     """
    103     logging.config.dictConfig({
    104         'version': 1,
    105         'formatters': {
    106             'stderr': {
    107                 'format': ('{name}: '
    108                            '%(asctime)s:%(levelname)s'
    109                            ':%(module)s:%(funcName)s:%(lineno)d'
    110                            ': %(message)s'
    111                            .format(name=name)),
    112             },
    113             'tee_file': {
    114                 'format': ('%(asctime)s:%(levelname)s'
    115                            ':%(module)s:%(funcName)s:%(lineno)d'
    116                            ': %(message)s'),
    117             },
    118         },
    119         'handlers': {
    120             'stderr': {
    121                 'class': 'logging.StreamHandler',
    122                 'formatter': 'stderr',
    123             },
    124             'tee_file': {
    125                 'class': 'logging.FileHandler',
    126                 'formatter': 'tee_file',
    127                 'filename': tee_file,
    128             },
    129         },
    130         'root': {
    131             'level': 'DEBUG',
    132             'handlers': ['stderr', 'tee_file'],
    133         },
    134         'disable_existing_loggers': False,
    135     })
    136 
    137 
    138 def _read_store(path):
    139   """Read a HostInfo from a file at path."""
    140   store = file_store.FileStore(path)
    141   return store.get()
    142 
    143 
    144 def _create_host(hostname, info, results_dir):
    145   """Yield a hosts.CrosHost object with the given inventory information.
    146 
    147   @param hostname: Hostname of the DUT.
    148   @param info: A HostInfo with the inventory information to use.
    149   @param results_dir: Path to directory for logs / output artifacts.
    150   @yield server.hosts.CrosHost object.
    151   """
    152   if not info.board:
    153     raise DutPreparationError('No board in DUT labels')
    154   if not info.model:
    155     raise DutPreparationError('No model in DUT labels')
    156 
    157   servo_args = {}
    158   if 'servo_host' not in info.attributes:
    159     raise DutPreparationError('No servo_host in DUT attributes')
    160   if 'servo_port' not in info.attributes:
    161     raise DutPreparationError('No servo_port in DUT attributes')
    162 
    163   dut_logs_dir = os.path.join(results_dir, _DUT_LOGS_DIR)
    164   try:
    165     os.makedirs(dut_logs_dir)
    166   except OSError as e:
    167     if e.errno != errno.EEXIST:
    168       raise
    169 
    170   return preparedut.create_host(
    171       hostname,
    172       info.board,
    173       info.model,
    174       info.attributes['servo_host'],
    175       info.attributes['servo_port'],
    176       info.attributes.get('servo_serial', ''),
    177       dut_logs_dir,
    178   )
    179 
    180 
    181 def _get_cros_repair_image_name(board):
    182   """Get the CrOS repair image name for given host.
    183 
    184   TODO(pprabhu): This is an evil function with dependence on the environment
    185   (global_config information) and the AFE. Remove this dependence when stable
    186   image mappings move off of the AFE.
    187   """
    188   return afe_utils.get_stable_cros_image_name(board)
    189 
    190 
    191 if __name__ == '__main__':
    192   main()
    193