1#!/usr/bin/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(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
54def 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
281def 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
317def 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
330def 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
345def 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
363def 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.
391def 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
408def 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
428if __name__ == '__main__':
429  retval = DoImage(sys.argv)
430  sys.exit(retval)
431