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