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