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