image_chromeos.py revision f395c26437cbdabc2960447fba89b226f4409e82
1#!/usr/bin/python2.6
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 fcntl
13import filecmp
14import glob
15import optparse
16import os
17import re
18import shutil
19import sys
20import tempfile
21from utils import command_executer
22from utils import logger
23from utils import misc
24from utils.file_utils import FileUtils
25
26checksum_file = "/usr/local/osimage_checksum_file"
27lock_file = "/tmp/lock_image_chromeos"
28
29def Usage(parser, message):
30  print "ERROR: " + message
31  parser.print_help()
32  sys.exit(0)
33
34
35def Main(argv):
36  """Build ChromeOS."""
37  #Get lock for the host
38  f = open(lock_file, "w+a")
39  try:
40    fcntl.lockf(f, fcntl.LOCK_EX|fcntl.LOCK_NB)
41  except IOError:
42    f.close()
43    print ("You can not run two instances of image_chromes at the same time."
44           "\nTry again. Exiting ....")
45    exit(0)
46  # Common initializations
47  cmd_executer = command_executer.GetCommandExecuter()
48  l = logger.GetLogger()
49
50  parser = optparse.OptionParser()
51  parser.add_option("-c", "--chromeos_root", dest="chromeos_root",
52                    help="Target directory for ChromeOS installation.")
53  parser.add_option("-r", "--remote", dest="remote",
54                    help="Target device.")
55  parser.add_option("-i", "--image", dest="image",
56                    help="Image binary file.")
57  parser.add_option("-b", "--board", dest="board",
58                    help="Target board override.")
59  parser.add_option("-f", "--force", dest="force",
60                    action="store_true",
61                    default=False,
62                    help="Force an image even if it is non-test.")
63  parser.add_option("-a",
64                    "--image_to_live_args",
65                    dest="image_to_live_args")
66
67
68  options = parser.parse_args(argv[1:])[0]
69
70  if options.chromeos_root is None:
71    Usage(parser, "--chromeos_root must be set")
72
73  if options.remote is None:
74    Usage(parser, "--remote must be set")
75
76  options.chromeos_root = os.path.expanduser(options.chromeos_root)
77
78  if options.board is None:
79    board = cmd_executer.CrosLearnBoard(options.chromeos_root, options.remote)
80  else:
81    board = options.board
82
83  if options.image is None:
84    images_dir = misc.GetImageDir(options.chromeos_root, board)
85    image = os.path.join(images_dir,
86                         "latest",
87                         "chromiumos_test_image.bin")
88    if not os.path.exists(image):
89      image = os.path.join(images_dir,
90                           "latest",
91                           "chromiumos_image.bin")
92  else:
93    image = options.image
94    image = os.path.expanduser(image)
95
96  image = os.path.realpath(image)
97
98  if not os.path.exists(image):
99    Usage(parser, "Image file: " + image + " does not exist!")
100
101  image_checksum = FileUtils().Md5File(image)
102
103  command = "cat " + checksum_file
104  retval, device_checksum, err = cmd_executer.CrosRunCommand(command,
105      return_output=True,
106      chromeos_root=options.chromeos_root,
107      machine=options.remote)
108
109  device_checksum = device_checksum.strip()
110  image_checksum = str(image_checksum)
111
112  l.LogOutput("Image checksum: " + image_checksum)
113  l.LogOutput("Device checksum: " + device_checksum)
114
115  if image_checksum != device_checksum:
116    [found, located_image] = LocateOrCopyImage(options.chromeos_root,
117                                               image,
118                                               board=board)
119
120    l.LogOutput("Checksums do not match. Re-imaging...")
121
122    is_test_image = IsImageModdedForTest(options.chromeos_root,
123                                         located_image)
124
125    if not is_test_image and not options.force:
126      logger.GetLogger().LogFatal("Have to pass --force to image a non-test "
127                                  "image!")
128
129    # If the device has /tmp mounted as noexec, image_to_live.sh can fail.
130    command = "mount -o remount,rw,exec /tmp"
131    cmd_executer.CrosRunCommand(command,
132                                chromeos_root=options.chromeos_root,
133                                machine=options.remote)
134
135    command = (options.chromeos_root +
136               "/src/scripts/image_to_live.sh --remote=" +
137               options.remote +
138               " --image=" + located_image)
139    if options.image_to_live_args:
140      command += " %s" % options.image_to_live_args
141
142    retval = cmd_executer.RunCommand(command)
143
144    if found == False:
145      temp_dir = os.path.dirname(located_image)
146      l.LogOutput("Deleting temp image dir: %s" % temp_dir)
147      shutil.rmtree(temp_dir)
148
149    logger.GetLogger().LogFatalIf(retval, "Image command failed")
150    command = "echo %s > %s && chmod -w %s" % (image_checksum, checksum_file,
151                                               checksum_file)
152    retval = cmd_executer.CrosRunCommand(command,
153                                         chromeos_root=options.chromeos_root,
154                                         machine=options.remote)
155    logger.GetLogger().LogFatalIf(retval, "Writing checksum failed.")
156
157    successfully_imaged = VerifyChromeChecksum(options.chromeos_root,
158                                               image,
159                                               options.remote)
160    logger.GetLogger().LogFatalIf(not successfully_imaged,
161                                  "Image verification failed!")
162  else:
163    l.LogOutput("Checksums match. Skipping reimage")
164
165  fcntl.lockf(f, fcntl.LOCK_UN)
166  f.close()
167  return retval
168
169
170def LocateOrCopyImage(chromeos_root, image, board=None):
171  l = logger.GetLogger()
172  if board is None:
173    board_glob = "*"
174  else:
175    board_glob = board
176
177  chromeos_root_realpath = os.path.realpath(chromeos_root)
178  image = os.path.realpath(image)
179
180  if image.startswith("%s/" % chromeos_root_realpath):
181    return [True, image]
182
183  # First search within the existing build dirs for any matching files.
184  images_glob = ("%s/src/build/images/%s/*/*.bin" %
185                 (chromeos_root_realpath,
186                  board_glob))
187  images_list = glob.glob(images_glob)
188  for potential_image in images_list:
189    if filecmp.cmp(potential_image, image):
190      l.LogOutput("Found matching image %s in chromeos_root." % potential_image)
191      return [True, potential_image]
192  # We did not find an image. Copy it in the src dir and return the copied file.
193  if board is None:
194    board = ""
195  base_dir = ("%s/src/build/images/%s" %
196              (chromeos_root_realpath,
197               board))
198  if not os.path.isdir(base_dir):
199    os.makedirs(base_dir)
200  temp_dir = tempfile.mkdtemp(prefix="%s/tmp" % base_dir)
201  new_image = "%s/%s" % (temp_dir, os.path.basename(image))
202  l.LogOutput("No matching image found. Copying %s to %s" %
203              (image, new_image))
204  shutil.copyfile(image, new_image)
205  return [False, new_image]
206
207
208def GetImageMountCommand(chromeos_root, image, rootfs_mp, stateful_mp):
209  image_dir = os.path.dirname(image)
210  image_file = os.path.basename(image)
211  mount_command = ("cd %s/src/scripts &&"
212                   "./mount_gpt_image.sh --from=%s --image=%s"
213                   " --safe --read_only"
214                   " --rootfs_mountpt=%s"
215                   " --stateful_mountpt=%s" %
216                   (chromeos_root, image_dir, image_file, rootfs_mp,
217                    stateful_mp))
218  return mount_command
219
220
221def MountImage(chromeos_root, image, rootfs_mp, stateful_mp, unmount=False):
222  cmd_executer = command_executer.GetCommandExecuter()
223  command = GetImageMountCommand(chromeos_root, image, rootfs_mp, stateful_mp)
224  if unmount:
225    command = "%s --unmount" % command
226  retval = cmd_executer.RunCommand(command)
227  logger.GetLogger().LogFatalIf(retval, "Mount/unmount command failed!")
228  return retval
229
230
231def IsImageModdedForTest(chromeos_root, image):
232  rootfs_mp = tempfile.mkdtemp()
233  stateful_mp = tempfile.mkdtemp()
234  MountImage(chromeos_root, image, rootfs_mp, stateful_mp)
235  lsb_release_file = os.path.join(rootfs_mp, "etc/lsb-release")
236  lsb_release_contents = open(lsb_release_file).read()
237  is_test_image = re.search("test", lsb_release_contents, re.IGNORECASE)
238  MountImage(chromeos_root, image, rootfs_mp, stateful_mp, unmount=True)
239  return is_test_image
240
241
242def VerifyChromeChecksum(chromeos_root, image, remote):
243  cmd_executer = command_executer.GetCommandExecuter()
244  rootfs_mp = tempfile.mkdtemp()
245  stateful_mp = tempfile.mkdtemp()
246  MountImage(chromeos_root, image, rootfs_mp, stateful_mp)
247  image_chrome_checksum = FileUtils().Md5File("%s/opt/google/chrome/chrome" %
248                                              rootfs_mp)
249  MountImage(chromeos_root, image, rootfs_mp, stateful_mp, unmount=True)
250
251  command = "md5sum /opt/google/chrome/chrome"
252  [r, o, e] = cmd_executer.CrosRunCommand(command,
253                                          return_output=True,
254                                          chromeos_root=chromeos_root,
255                                          machine=remote)
256  device_chrome_checksum = o.split()[0]
257  if image_chrome_checksum.strip() == device_chrome_checksum.strip():
258    return True
259  else:
260    return False
261
262
263if __name__ == "__main__":
264  retval = Main(sys.argv)
265  sys.exit(retval)
266