• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/python2 -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
8from __future__ import absolute_import
9from __future__ import division
10from __future__ import print_function
11
12import argparse
13import errno
14import logging
15import logging.config
16import os
17import sys
18
19import common
20from autotest_lib.client.common_lib import enum
21from autotest_lib.server import afe_utils
22from autotest_lib.server.hosts import file_store
23from autotest_lib.site_utils.deployment.prepare import dut as preparedut
24
25
26RETURN_CODES = enum.Enum(
27        'OK',
28        'STAGE_USB_FAILURE',
29        'INSTALL_FIRMWARE_FAILURE',
30        'INSTALL_TEST_IMAGE_FAILURE',
31        'BOOT_FROM_RECOVERY_MODE_FAILURE',
32        'SETUP_LABSTATION_FAILURE',
33        'UPDATE_LABEL_FAILURE',
34        'OTHER_FAILURES',
35)
36
37
38class DutPreparationError(Exception):
39  """Generic error raised during DUT preparation."""
40
41
42def main():
43  """Tool to (re)prepare a DUT for lab deployment."""
44  opts = _parse_args()
45  _configure_logging('prepare_dut', os.path.join(opts.results_dir, _LOG_FILE))
46
47  try:
48    info = _read_store(opts.host_info_file)
49  except Exception as err:
50    logging.error("fail to prepare: %s", err)
51    return RETURN_CODES.OTHER_FAILURES
52
53  with _create_host(opts.hostname, info, opts.results_dir) as host:
54    if opts.dry_run:
55      logging.info('DRY RUN: Would have run actions %s', opts.actions)
56      return
57
58    if 'stage-usb' in opts.actions:
59      try:
60        repair_image = afe_utils.get_stable_cros_image_name_v2(info)
61        logging.info('Using repair image %s, obtained from AFE', repair_image)
62        preparedut.download_image_to_servo_usb(host, repair_image)
63      except Exception as err:
64        logging.error("fail to stage image to usb: %s", err)
65        return RETURN_CODES.STAGE_USB_FAILURE
66
67    if 'install-test-image' in opts.actions:
68      try:
69        preparedut.install_test_image(host)
70      except Exception as err:
71        logging.error("fail to install test image: %s", err)
72        return RETURN_CODES.INSTALL_TEST_IMAGE_FAILURE
73
74    if 'install-firmware' in opts.actions:
75      try:
76        preparedut.install_firmware(host)
77      except Exception as err:
78        logging.error("fail to install firmware: %s", err)
79        return RETURN_CODES.INSTALL_FIRMWARE_FAILURE
80
81    if 'verify-recovery-mode' in opts.actions:
82      try:
83        preparedut.verify_boot_into_rec_mode(host)
84      except Exception as err:
85        logging.error("fail to boot from recovery mode: %s", err)
86        return RETURN_CODES.BOOT_FROM_RECOVERY_MODE_FAILURE
87
88    if 'setup-labstation' in opts.actions:
89      try:
90        preparedut.setup_labstation(host)
91      except Exception as err:
92        logging.error("fail to setup labstation: %s", err)
93        return RETURN_CODES.SETUP_LABSTATION_FAILURE
94
95    if 'update-label' in opts.actions:
96      try:
97        host.labels.update_labels(host, task_name='deploy')
98      except Exception as err:
99        logging.error("fail to update label: %s", err)
100        return RETURN_CODES.UPDATE_LABEL_FAILURE
101
102  return RETURN_CODES.OK
103
104
105_LOG_FILE = 'prepare_dut.log'
106_DUT_LOGS_DIR = 'dut_logs'
107
108
109def _parse_args():
110  parser = argparse.ArgumentParser(
111      description='Prepare / validate DUT for lab deployment.')
112
113  parser.add_argument(
114      'actions',
115      nargs='+',
116      choices=['stage-usb', 'install-test-image', 'install-firmware',
117               'verify-recovery-mode', 'update-label', 'setup-labstation'],
118      help='DUT preparation actions to execute.',
119  )
120  parser.add_argument(
121      '--dry-run',
122      action='store_true',
123      default=False,
124      help='Run in dry-run mode. No changes will be made to the DUT.',
125  )
126  parser.add_argument(
127      '--results-dir',
128      required=True,
129      help='Directory to drop logs and output artifacts in.',
130  )
131
132  parser.add_argument(
133      '--hostname',
134      required=True,
135      help='Hostname of the DUT to prepare.',
136  )
137  parser.add_argument(
138      '--host-info-file',
139      required=True,
140      help=('Full path to HostInfo file.'
141            ' DUT inventory information is read from the HostInfo file.'),
142  )
143
144  return parser.parse_args()
145
146
147def _configure_logging(name, tee_file):
148    """Configure logging globally.
149
150    @param name: Name to prepend to log messages.
151                 This should be the name of the program.
152    @param tee_file: File to tee logs to, in addition to stderr.
153    """
154    logging.config.dictConfig({
155        'version': 1,
156        'formatters': {
157            'stderr': {
158                'format': ('{name}: '
159                           '%(asctime)s:%(levelname)s'
160                           ':%(module)s:%(funcName)s:%(lineno)d'
161                           ': %(message)s'
162                           .format(name=name)),
163            },
164            'tee_file': {
165                'format': ('%(asctime)s:%(levelname)s'
166                           ':%(module)s:%(funcName)s:%(lineno)d'
167                           ': %(message)s'),
168            },
169        },
170        'handlers': {
171            'stderr': {
172                'class': 'logging.StreamHandler',
173                'formatter': 'stderr',
174            },
175            'tee_file': {
176                'class': 'logging.FileHandler',
177                'formatter': 'tee_file',
178                'filename': tee_file,
179            },
180        },
181        'root': {
182            'level': 'DEBUG',
183            'handlers': ['stderr', 'tee_file'],
184        },
185        'disable_existing_loggers': False,
186    })
187
188
189def _read_store(path):
190  """Read a HostInfo from a file at path."""
191  store = file_store.FileStore(path)
192  return store.get()
193
194
195def _create_host(hostname, info, results_dir):
196  """Yield a hosts.CrosHost object with the given inventory information.
197
198  @param hostname: Hostname of the DUT.
199  @param info: A HostInfo with the inventory information to use.
200  @param results_dir: Path to directory for logs / output artifacts.
201  @yield server.hosts.CrosHost object.
202  """
203  if not info.board:
204    raise DutPreparationError('No board in DUT labels')
205  if not info.model:
206    raise DutPreparationError('No model in DUT labels')
207
208  if info.os == 'labstation':
209    return preparedut.create_labstation_host(hostname, info.board, info.model)
210
211  # We assume target host is a cros DUT by default
212  if 'servo_host' not in info.attributes:
213    raise DutPreparationError('No servo_host in DUT attributes')
214  if 'servo_port' not in info.attributes:
215    raise DutPreparationError('No servo_port in DUT attributes')
216
217  dut_logs_dir = os.path.join(results_dir, _DUT_LOGS_DIR)
218  try:
219    os.makedirs(dut_logs_dir)
220  except OSError as e:
221    if e.errno != errno.EEXIST:
222      raise
223
224  return preparedut.create_cros_host(
225      hostname,
226      info.board,
227      info.model,
228      info.attributes['servo_host'],
229      info.attributes['servo_port'],
230      info.attributes.get('servo_serial'),
231      dut_logs_dir,
232  )
233
234
235if __name__ == '__main__':
236  sys.exit(main())
237