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