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