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