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