• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python2
2#
3# Copyright 2011 Google Inc. All Rights Reserved.
4"""Script to image a ChromeOS device.
5
6This script images a remote ChromeOS device with a specific image."
7"""
8
9from __future__ import print_function
10
11__author__ = 'asharif@google.com (Ahmad Sharif)'
12
13import argparse
14import filecmp
15import glob
16import os
17import re
18import shutil
19import sys
20import tempfile
21import time
22
23from cros_utils import command_executer
24from cros_utils import locks
25from cros_utils import logger
26from cros_utils import misc
27from cros_utils.file_utils import FileUtils
28
29checksum_file = '/usr/local/osimage_checksum_file'
30lock_file = '/tmp/image_chromeos_lock/image_chromeos_lock'
31
32
33def Usage(parser, message):
34  print('ERROR: %s' % message)
35  parser.print_help()
36  sys.exit(0)
37
38
39def CheckForCrosFlash(chromeos_root, remote, log_level):
40  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
41
42  # Check to see if remote machine has cherrypy, ctypes
43  command = "python -c 'import cherrypy, ctypes'"
44  ret = cmd_executer.CrosRunCommand(
45      command, chromeos_root=chromeos_root, machine=remote)
46  logger.GetLogger().LogFatalIf(
47      ret == 255, 'Failed ssh to %s (for checking cherrypy)' % remote)
48  logger.GetLogger().LogFatalIf(
49      ret != 0, "Failed to find cherrypy or ctypes on remote '{}', "
50      'cros flash cannot work.'.format(remote))
51
52
53def DisableCrosBeeps(chromeos_root, remote, log_level):
54  """Disable annoying chromebooks beeps after reboots."""
55  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
56
57  command = '/usr/share/vboot/bin/set_gbb_flags.sh 0x1'
58  logger.GetLogger().LogOutput('Trying to disable beeping.')
59
60  ret, o, _ = cmd_executer.CrosRunCommandWOutput(
61      command, chromeos_root=chromeos_root, machine=remote)
62  if ret != 0:
63    logger.GetLogger().LogOutput(o)
64    logger.GetLogger().LogOutput('Failed to disable beeps.')
65
66
67def DoImage(argv):
68  """Image ChromeOS."""
69
70  parser = argparse.ArgumentParser()
71  parser.add_argument(
72      '-c',
73      '--chromeos_root',
74      dest='chromeos_root',
75      help='Target directory for ChromeOS installation.')
76  parser.add_argument('-r', '--remote', dest='remote', help='Target device.')
77  parser.add_argument('-i', '--image', dest='image', help='Image binary file.')
78  parser.add_argument(
79      '-b', '--board', dest='board', help='Target board override.')
80  parser.add_argument(
81      '-f',
82      '--force',
83      dest='force',
84      action='store_true',
85      default=False,
86      help='Force an image even if it is non-test.')
87  parser.add_argument(
88      '-n',
89      '--no_lock',
90      dest='no_lock',
91      default=False,
92      action='store_true',
93      help='Do not attempt to lock remote before imaging.  '
94      'This option should only be used in cases where the '
95      'exclusive lock has already been acquired (e.g. in '
96      'a script that calls this one).')
97  parser.add_argument(
98      '-l',
99      '--logging_level',
100      dest='log_level',
101      default='verbose',
102      help='Amount of logging to be used. Valid levels are '
103      "'quiet', 'average', and 'verbose'.")
104  parser.add_argument('-a', '--image_args', dest='image_args')
105
106  options = parser.parse_args(argv[1:])
107
108  if not options.log_level in command_executer.LOG_LEVEL:
109    Usage(parser, "--logging_level must be 'quiet', 'average' or 'verbose'")
110  else:
111    log_level = options.log_level
112
113  # Common initializations
114  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
115  l = logger.GetLogger()
116
117  if options.chromeos_root is None:
118    Usage(parser, '--chromeos_root must be set')
119
120  if options.remote is None:
121    Usage(parser, '--remote must be set')
122
123  options.chromeos_root = os.path.expanduser(options.chromeos_root)
124
125  if options.board is None:
126    board = cmd_executer.CrosLearnBoard(options.chromeos_root, options.remote)
127  else:
128    board = options.board
129
130  if options.image is None:
131    images_dir = misc.GetImageDir(options.chromeos_root, board)
132    image = os.path.join(images_dir, 'latest', 'chromiumos_test_image.bin')
133    if not os.path.exists(image):
134      image = os.path.join(images_dir, 'latest', 'chromiumos_image.bin')
135    is_xbuddy_image = False
136  else:
137    image = options.image
138    is_xbuddy_image = image.startswith('xbuddy://')
139    if not is_xbuddy_image:
140      image = os.path.expanduser(image)
141
142  if not is_xbuddy_image:
143    image = os.path.realpath(image)
144
145  if not os.path.exists(image) and not is_xbuddy_image:
146    Usage(parser, 'Image file: ' + image + ' does not exist!')
147
148  try:
149    should_unlock = False
150    if not options.no_lock:
151      try:
152        _ = locks.AcquireLock(
153            list(options.remote.split()), options.chromeos_root)
154        should_unlock = True
155      except Exception as e:
156        raise RuntimeError('Error acquiring machine: %s' % str(e))
157
158    reimage = False
159    local_image = False
160    if not is_xbuddy_image:
161      local_image = True
162      image_checksum = FileUtils().Md5File(image, log_level=log_level)
163
164      command = 'cat ' + checksum_file
165      ret, device_checksum, _ = cmd_executer.CrosRunCommandWOutput(
166          command, chromeos_root=options.chromeos_root, machine=options.remote)
167
168      device_checksum = device_checksum.strip()
169      image_checksum = str(image_checksum)
170
171      l.LogOutput('Image checksum: ' + image_checksum)
172      l.LogOutput('Device checksum: ' + device_checksum)
173
174      if image_checksum != device_checksum:
175        [found, located_image] = LocateOrCopyImage(
176            options.chromeos_root, image, board=board)
177
178        reimage = True
179        l.LogOutput('Checksums do not match. Re-imaging...')
180
181        is_test_image = IsImageModdedForTest(options.chromeos_root,
182                                             located_image, log_level)
183
184        if not is_test_image and not options.force:
185          logger.GetLogger().LogFatal('Have to pass --force to image a '
186                                      'non-test image!')
187    else:
188      reimage = True
189      found = True
190      l.LogOutput('Using non-local image; Re-imaging...')
191
192    if reimage:
193      # If the device has /tmp mounted as noexec, image_to_live.sh can fail.
194      command = 'mount -o remount,rw,exec /tmp'
195      cmd_executer.CrosRunCommand(
196          command, chromeos_root=options.chromeos_root, machine=options.remote)
197
198      real_src_dir = os.path.join(
199          os.path.realpath(options.chromeos_root), 'src')
200      real_chroot_dir = os.path.join(
201          os.path.realpath(options.chromeos_root), 'chroot')
202      if local_image:
203        if located_image.find(real_src_dir) != 0:
204          if located_image.find(real_chroot_dir) != 0:
205            raise RuntimeError('Located image: %s not in chromeos_root: %s' %
206                               (located_image, options.chromeos_root))
207          else:
208            chroot_image = located_image[len(real_chroot_dir):]
209        else:
210          chroot_image = os.path.join(
211              '~/trunk/src', located_image[len(real_src_dir):].lstrip('/'))
212
213      # Check to see if cros flash will work for the remote machine.
214      CheckForCrosFlash(options.chromeos_root, options.remote, log_level)
215
216      # Disable the annoying chromebook beeps after reboot.
217      DisableCrosBeeps(options.chromeos_root, options.remote, log_level)
218
219      cros_flash_args = [
220          'cros', 'flash',
221          '--board=%s' % board, '--clobber-stateful', options.remote
222      ]
223      if local_image:
224        cros_flash_args.append(chroot_image)
225      else:
226        cros_flash_args.append(image)
227
228      command = ' '.join(cros_flash_args)
229
230      # Workaround for crosbug.com/35684.
231      os.chmod(misc.GetChromeOSKeyFile(options.chromeos_root), 0600)
232
233      if log_level == 'average':
234        cmd_executer.SetLogLevel('verbose')
235      retries = 0
236      while True:
237        if log_level == 'quiet':
238          l.LogOutput('CMD : %s' % command)
239        ret = cmd_executer.ChrootRunCommand(
240            options.chromeos_root, command, command_timeout=1800)
241        if ret == 0 or retries >= 2:
242          break
243        retries += 1
244        if log_level == 'quiet':
245          l.LogOutput('Imaging failed. Retry # %d.' % retries)
246
247      if log_level == 'average':
248        cmd_executer.SetLogLevel(log_level)
249
250      if found == False:
251        temp_dir = os.path.dirname(located_image)
252        l.LogOutput('Deleting temp image dir: %s' % temp_dir)
253        shutil.rmtree(temp_dir)
254
255      logger.GetLogger().LogFatalIf(ret, 'Image command failed')
256
257      # Unfortunately cros_image_to_target.py sometimes returns early when the
258      # machine isn't fully up yet.
259      ret = EnsureMachineUp(options.chromeos_root, options.remote, log_level)
260
261      # If this is a non-local image, then the ret returned from
262      # EnsureMachineUp is the one that will be returned by this function;
263      # in that case, make sure the value in 'ret' is appropriate.
264      if not local_image and ret == True:
265        ret = 0
266      else:
267        ret = 1
268
269      if local_image:
270        if log_level == 'average':
271          l.LogOutput('Verifying image.')
272        command = 'echo %s > %s && chmod -w %s' % (image_checksum,
273                                                   checksum_file, checksum_file)
274        ret = cmd_executer.CrosRunCommand(
275            command,
276            chromeos_root=options.chromeos_root,
277            machine=options.remote)
278        logger.GetLogger().LogFatalIf(ret, 'Writing checksum failed.')
279
280        successfully_imaged = VerifyChromeChecksum(options.chromeos_root, image,
281                                                   options.remote, log_level)
282        logger.GetLogger().LogFatalIf(not successfully_imaged,
283                                      'Image verification failed!')
284        TryRemountPartitionAsRW(options.chromeos_root, options.remote,
285                                log_level)
286    else:
287      l.LogOutput('Checksums match. Skipping reimage')
288    return ret
289  finally:
290    if should_unlock:
291      locks.ReleaseLock(list(options.remote.split()), options.chromeos_root)
292
293
294def LocateOrCopyImage(chromeos_root, image, board=None):
295  l = logger.GetLogger()
296  if board is None:
297    board_glob = '*'
298  else:
299    board_glob = board
300
301  chromeos_root_realpath = os.path.realpath(chromeos_root)
302  image = os.path.realpath(image)
303
304  if image.startswith('%s/' % chromeos_root_realpath):
305    return [True, image]
306
307  # First search within the existing build dirs for any matching files.
308  images_glob = ('%s/src/build/images/%s/*/*.bin' % (chromeos_root_realpath,
309                                                     board_glob))
310  images_list = glob.glob(images_glob)
311  for potential_image in images_list:
312    if filecmp.cmp(potential_image, image):
313      l.LogOutput('Found matching image %s in chromeos_root.' % potential_image)
314      return [True, potential_image]
315  # We did not find an image. Copy it in the src dir and return the copied
316  # file.
317  if board is None:
318    board = ''
319  base_dir = ('%s/src/build/images/%s' % (chromeos_root_realpath, board))
320  if not os.path.isdir(base_dir):
321    os.makedirs(base_dir)
322  temp_dir = tempfile.mkdtemp(prefix='%s/tmp' % base_dir)
323  new_image = '%s/%s' % (temp_dir, os.path.basename(image))
324  l.LogOutput('No matching image found. Copying %s to %s' % (image, new_image))
325  shutil.copyfile(image, new_image)
326  return [False, new_image]
327
328
329def GetImageMountCommand(chromeos_root, image, rootfs_mp, stateful_mp):
330  image_dir = os.path.dirname(image)
331  image_file = os.path.basename(image)
332  mount_command = ('cd %s/src/scripts &&'
333                   './mount_gpt_image.sh --from=%s --image=%s'
334                   ' --safe --read_only'
335                   ' --rootfs_mountpt=%s'
336                   ' --stateful_mountpt=%s' %
337                   (chromeos_root, image_dir, image_file, rootfs_mp,
338                    stateful_mp))
339  return mount_command
340
341
342def MountImage(chromeos_root,
343               image,
344               rootfs_mp,
345               stateful_mp,
346               log_level,
347               unmount=False):
348  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
349  command = GetImageMountCommand(chromeos_root, image, rootfs_mp, stateful_mp)
350  if unmount:
351    command = '%s --unmount' % command
352  ret = cmd_executer.RunCommand(command)
353  logger.GetLogger().LogFatalIf(ret, 'Mount/unmount command failed!')
354  return ret
355
356
357def IsImageModdedForTest(chromeos_root, image, log_level):
358  if log_level != 'verbose':
359    log_level = 'quiet'
360  rootfs_mp = tempfile.mkdtemp()
361  stateful_mp = tempfile.mkdtemp()
362  MountImage(chromeos_root, image, rootfs_mp, stateful_mp, log_level)
363  lsb_release_file = os.path.join(rootfs_mp, 'etc/lsb-release')
364  lsb_release_contents = open(lsb_release_file).read()
365  is_test_image = re.search('test', lsb_release_contents, re.IGNORECASE)
366  MountImage(
367      chromeos_root, image, rootfs_mp, stateful_mp, log_level, unmount=True)
368  return is_test_image
369
370
371def VerifyChromeChecksum(chromeos_root, image, remote, log_level):
372  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
373  rootfs_mp = tempfile.mkdtemp()
374  stateful_mp = tempfile.mkdtemp()
375  MountImage(chromeos_root, image, rootfs_mp, stateful_mp, log_level)
376  image_chrome_checksum = FileUtils().Md5File(
377      '%s/opt/google/chrome/chrome' % rootfs_mp, log_level=log_level)
378  MountImage(
379      chromeos_root, image, rootfs_mp, stateful_mp, log_level, unmount=True)
380
381  command = 'md5sum /opt/google/chrome/chrome'
382  [_, o, _] = cmd_executer.CrosRunCommandWOutput(
383      command, chromeos_root=chromeos_root, machine=remote)
384  device_chrome_checksum = o.split()[0]
385  if image_chrome_checksum.strip() == device_chrome_checksum.strip():
386    return True
387  else:
388    return False
389
390
391# Remount partition as writable.
392# TODO: auto-detect if an image is built using --noenable_rootfs_verification.
393def TryRemountPartitionAsRW(chromeos_root, remote, log_level):
394  l = logger.GetLogger()
395  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
396  command = 'sudo mount -o remount,rw /'
397  ret = cmd_executer.CrosRunCommand(\
398    command, chromeos_root=chromeos_root, machine=remote,
399    terminated_timeout=10)
400  if ret:
401    ## Safely ignore.
402    l.LogWarning('Failed to remount partition as rw, '
403                 'probably the image was not built with '
404                 "\"--noenable_rootfs_verification\", "
405                 'you can safely ignore this.')
406  else:
407    l.LogOutput('Re-mounted partition as writable.')
408
409
410def EnsureMachineUp(chromeos_root, remote, log_level):
411  l = logger.GetLogger()
412  cmd_executer = command_executer.GetCommandExecuter(log_level=log_level)
413  timeout = 600
414  magic = 'abcdefghijklmnopqrstuvwxyz'
415  command = 'echo %s' % magic
416  start_time = time.time()
417  while True:
418    current_time = time.time()
419    if current_time - start_time > timeout:
420      l.LogError(
421          'Timeout of %ss reached. Machine still not up. Aborting.' % timeout)
422      return False
423    ret = cmd_executer.CrosRunCommand(
424        command, chromeos_root=chromeos_root, machine=remote)
425    if not ret:
426      return True
427
428
429if __name__ == '__main__':
430  retval = DoImage(sys.argv)
431  sys.exit(retval)
432