src/third_party/chromite/cros/commands/cros_flash.py

   1 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Install/copy the image to the device."""
   6
   7 import cStringIO
   8 import logging
   9 import os
  10 import shutil
  11 import tempfile
  12 import time
  13 import urlparse
  14
  15 from chromite import cros
  16 from chromite.buildbot import constants
  17 from chromite.lib import cros_build_lib
  18 from chromite.lib import dev_server_wrapper as ds_wrapper
  19 from chromite.lib import osutils
  20 from chromite.lib import remote_access
  21
  22
  23 DEVSERVER_STATIC_DIR = cros_build_lib.FromChrootPath(
  24     os.path.join(constants.CHROOT_SOURCE_ROOT, 'devserver', 'static'))
  25
  26 IMAGE_NAME_TO_TYPE = {
  27     'chromiumos_test_image.bin': 'test',
  28     'chromiumos_image.bin': 'dev',
  29     'chromiumos_base_image.bin': 'base',
  30     'recovery_image.bin': 'recovery',
  31 }
  32
  33 IMAGE_TYPE_TO_NAME = {
  34     'test': 'chromiumos_test_image.bin',
  35     'dev': 'chromiumos_image.bin',
  36     'base': 'chromiumos_base_image.bin',
  37     'recovery': 'recovery_image.bin',
  38 }
  39
  40 XBUDDY_REMOTE = 'remote'
  41 XBUDDY_LOCAL = 'local'
  42
  43
  44 def ConvertTranslatedPath(original_path, translated_path):
  45   """Converts a translated xbuddy path to an xbuddy path.
  46
  47   Devserver/xbuddy does not accept requests with translated xbuddy
  48   path (build-id/version/image-name). This function converts such a
  49   translated path to an xbuddy path that is suitable to used in
  50   devserver requests.
  51
  52   Args:
  53     original_path: the xbuddy path before translation.
  54       (e.g., remote/peppy/latest-canary).
  55     translated_path: the translated xbuddy path
  56       (e.g., peppy-release/R36-5760.0.0).
  57
  58   Returns:
  59     A xbuddy path uniquely identifies a build and can be used in devserver
  60       requests: {local|remote}/build-id/version/image_type
  61   """
  62   chunks = translated_path.split(os.path.sep)
  63   chunks[-1] = IMAGE_NAME_TO_TYPE[chunks[-1]]
  64
  65   if _GetXbuddyPath(original_path).startswith(XBUDDY_REMOTE):
  66     chunks = [XBUDDY_REMOTE] + chunks
  67   else:
  68     chunks = [XBUDDY_LOCAL] + chunks
  69
  70   return os.path.sep.join(chunks)
  71
  72
  73 def _GetXbuddyPath(path):
  74   """A helper function to parse an xbuddy path.
  75
  76   Args:
  77     path: Either a path without no scheme or an xbuddy://path/for/xbuddy
  78
  79   Returns:
  80     path/for/xbuddy if |path| is xbuddy://path/for/xbuddy; otherwise,
  81     returns |path|.
  82
  83   Raises:
  84     ValueError if |path| uses any scheme other than xbuddy://.
  85   """
  86   parsed = urlparse.urlparse(path)
  87
  88   # pylint: disable=E1101
  89   if parsed.scheme == 'xbuddy':
  90     return '%s%s' % (parsed.netloc, parsed.path)
  91   elif parsed.scheme == '':
  92     logging.debug('Assuming %s is an xbuddy path.', path)
  93     return path
  94   else:
  95     raise ValueError('Do not support scheme %s.', parsed.scheme)
  96
  97
  98 def TranslateImagePath(path, board):
  99   """Start devserver to translate the xbuddy |path|.
 100
 101   Args:
 102     path: The xbuddy path.
 103     board: The default board to use if board is not specified in |path|.
 104
 105   Returns:
 106     A translated path that uniquely identifies one build:
 107       build-id/version/image_name
 108   """
 109   ds = ds_wrapper.DevServerWrapper(static_dir=DEVSERVER_STATIC_DIR,
 110                                    board=board)
 111   req = GenerateXbuddyRequest(path, 'translate')
 112   logging.info('Starting local devserver to get image path...')
 113   try:
 114     ds.Start()
 115     return ds.OpenURL(ds.GetDevServerURL(sub_dir=req), timeout=60 * 15)
 116
 117   except ds_wrapper.DevServerResponseError:
 118     logging.error('Unable to translate the image path: %s. Are you sure the '
 119                   'image path is correct? The board %s is used when no board '
 120                   'name is included in the image path.', path, board)
 121     raise ValueError('Cannot locate image: %s' % path)
 122   except ds_wrapper.DevServerException:
 123     logging.warning(ds.TailLog() or 'No devserver log is available.')
 124     raise
 125   finally:
 126     ds.Stop()
 127
 128
 129 def GenerateXbuddyRequest(path, req_type):
 130   """Generate an xbuddy request used to retreive payloads.
 131
 132   This function generates a xbuddy request based on |path| and
 133   |req_type|, which can be used to query the devserver. For request
 134   type 'image' ('update'), the devserver will repond with a URL
 135   pointing to the folder where the image (update payloads) is stored.
 136
 137   Args:
 138     path: An xbuddy path (with or without xbuddy://).
 139     req_type: xbuddy request type ('update', 'image', or 'translate').
 140
 141   Returns:
 142     A xbuddy request.
 143   """
 144   if req_type == 'update':
 145     return 'xbuddy/%s?for_update=true&return_dir=true' % _GetXbuddyPath(path)
 146   elif req_type == 'image':
 147     return 'xbuddy/%s?return_dir=true' % _GetXbuddyPath(path)
 148   elif req_type == 'translate':
 149     return 'xbuddy_translate/%s' % _GetXbuddyPath(path)
 150   else:
 151     raise ValueError('Does not support xbuddy request type %s' % req_type)
 152
 153
 154 def DevserverURLToLocalPath(url, static_dir, file_type):
 155   """Convert the devserver returned URL to a local path.
 156
 157   Devserver returns only the directory where the files are. This
 158   function converts such a URL to a local path based on |file_type| so
 159   that we can access the file without downloading it.
 160
 161   Args:
 162     url: The URL returned by devserver (when return_dir=true).
 163     static_dir: The static directory used by the devserver.
 164     file_type: The image (in IMAGE_TYPE_TO_NAME) that we want to access.
 165
 166   Returns:
 167     A local path to the file.
 168   """
 169   # pylint: disable=E1101
 170   # Example URL: http://localhost:8080/static/peppy-release/R33-5116.87.0
 171   relative_path = urlparse.urlparse(url).path[len('/static/'):]
 172   # Defaults to test image because that is how Xbuddy handles the path.
 173   filename = IMAGE_TYPE_TO_NAME.get(file_type, IMAGE_TYPE_TO_NAME['test'])
 174   # Expand the path because devserver may use symlinks.
 175   real_path = osutils.ExpandPath(
 176       os.path.join(static_dir, relative_path, filename))
 177
 178   # If devserver uses a symlink within chroot, and we are running
 179   # outside of chroot, we need to convert the path.
 180   if os.path.exists(real_path):
 181     return real_path
 182   else:
 183     return cros_build_lib.FromChrootPath(real_path)
 184
 185
 186 class USBImager(object):
 187   """Copy image to the target removable device."""
 188
 189   def __init__(self, device, board, image, debug=False, yes=False):
 190     """Initalizes USBImager."""
 191     self.device = device
 192     self.board = board if board else cros_build_lib.GetDefaultBoard()
 193     self.image = image
 194     self.debug = debug
 195     self.debug_level = logging.DEBUG if debug else logging.INFO
 196     self.yes = yes
 197
 198   def DeviceNameToPath(self, device_name):
 199     return '/dev/%s' % device_name
 200
 201   def GetRemovableDeviceDescription(self, device):
 202     """Returns a informational description of the removable |device|.
 203
 204     Args:
 205       device: the device name (e.g. sdc).
 206
 207     Returns:
 208       A string describing |device| (e.g. Patriot Memory 7918 MB).
 209     """
 210     desc = []
 211     desc.append(osutils.GetDeviceInfo(device, keyword='manufacturer'))
 212     desc.append(osutils.GetDeviceInfo(device, keyword='product'))
 213     desc.append(osutils.GetDeviceSize(self.DeviceNameToPath(device)))
 214     return ' '.join([x for x in desc if x])
 215
 216   def ListAllRemovableDevices(self):
 217     """Returns a list of removable devices.
 218
 219     Returns:
 220       A list of device names (e.g. ['sdb', 'sdc']).
 221     """
 222     devices = osutils.ListBlockDevices()
 223     removable_devices = []
 224     for d in devices:
 225       if d.TYPE == 'disk' and d.RM == '1':
 226         removable_devices.append(d.NAME)
 227
 228     return removable_devices
 229
 230   def ChooseRemovableDevice(self, devices):
 231     """Lists all removable devices and asks user to select/confirm.
 232
 233     Args:
 234       devices: a list of device names (e.g. ['sda', 'sdb']).
 235
 236     Returns:
 237       The device name chosen by the user.
 238     """
 239     idx = cros_build_lib.GetChoice(
 240       'Removable device(s) found. Please select/confirm to continue:',
 241       [self.GetRemovableDeviceDescription(x) for x in devices])
 242
 243     return devices[idx]
 244
 245   def CopyImageToDevice(self, image, device):
 246     """Copies |image |to the removable |device|.
 247
 248     Args:
 249       image: Path to the image to copy.
 250       device: Device to copy to.
 251     """
 252     # Use pv to display progress bar if possible.
 253     cmd_base = 'pv -pretb'
 254     try:
 255       cros_build_lib.RunCommand(['pv', '--version'], print_cmd=False,
 256                                 capture_output=True)
 257     except cros_build_lib.RunCommandError:
 258       cmd_base = 'cat'
 259
 260     cmd = '%s %s | dd of=%s bs=4M oflag=sync' % (cmd_base, image, device)
 261     cros_build_lib.SudoRunCommand(cmd, shell=True)
 262     cros_build_lib.SudoRunCommand(['sync'], debug_level=self.debug_level)
 263
 264   def GetImagePathFromDevserver(self, path):
 265     """Gets image path from devserver.
 266
 267     Asks devserver to stage the image and convert the returned URL to a
 268     local path to the image.
 269
 270     Args:
 271       path: An xbuddy path with or without (xbuddy://).
 272
 273     Returns:
 274       A local path to the image.
 275     """
 276     ds = ds_wrapper.DevServerWrapper(static_dir=DEVSERVER_STATIC_DIR,
 277                                      board=self.board)
 278     req = GenerateXbuddyRequest(path, 'image')
 279     logging.info('Starting a local devserver to stage image...')
 280     try:
 281       ds.Start()
 282       url = ds.OpenURL(ds.GetDevServerURL(sub_dir=req), timeout=60 * 15)
 283
 284     except ds_wrapper.DevServerResponseError:
 285       logging.warning('Could not download %s.', path)
 286       logging.warning(ds.TailLog() or 'No devserver log is available.')
 287       raise
 288     else:
 289       # Print out the log when debug is on.
 290       logging.debug(ds.TailLog() or 'No devserver log is available.')
 291     finally:
 292       ds.Stop()
 293
 294     return DevserverURLToLocalPath(url, DEVSERVER_STATIC_DIR,
 295                                    path.rsplit(os.path.sep)[-1])
 296
 297   def ChooseImageFromDirectory(self, dir_path):
 298     """Lists all image files in |dir_path| and ask user to select one."""
 299     images = [x for x in os.listdir(dir_path) if
 300               os.path.isfile(os.path.join(dir_path, x)) and x.endswith(".bin")]
 301     idx = 0
 302     if len(images) == 0:
 303       raise ValueError('No image found in %s.' % dir_path)
 304     elif len(images) > 1:
 305       idx = cros_build_lib.GetChoice(
 306           'Multiple images found in %s. Please select one to continue:' % (
 307           dir_path), images)
 308
 309     return os.path.join(dir_path, images[idx])
 310
 311   def _GetImagePath(self):
 312     """Returns the image path to use."""
 313     image_path = translated_path = None
 314     if os.path.isfile(self.image):
 315       image_path = self.image
 316     elif os.path.isdir(self.image):
 317       # Ask user which image (*.bin) in the folder to use.
 318       image_path = self.ChooseImageFromDirectory(self.image)
 319     else:
 320       # Translate the xbuddy path to get the exact image to use.
 321       translated_path = TranslateImagePath(self.image, self.board)
 322       # Convert the translated path to be used in a request.
 323       xbuddy_path = ConvertTranslatedPath(self.image, translated_path)
 324       image_path = self.GetImagePathFromDevserver(xbuddy_path)
 325
 326     logging.info('Using image %s', translated_path or image_path)
 327     return image_path
 328
 329   def Run(self):
 330     """Image the removable device."""
 331     devices = self.ListAllRemovableDevices()
 332     # If user has specified a device, check if it is removable.
 333     if (self.device and
 334         self.device not in [self.DeviceNameToPath(x) for x in devices]):
 335       msg = '%s is not a removable device.' % self.device
 336       if not (self.yes or cros_build_lib.BooleanPrompt(
 337           default=False, prolog=msg)):
 338         cros_build_lib.Die('You can specify usb:// to choose from a list of '
 339                            'removable devices.')
 340     target = None
 341     if self.device:
 342       # Get device name from path (e.g. sdc in /dev/sdc).
 343       target = self.device.rsplit(os.path.sep, 1)[-1]
 344     elif devices:
 345       # Ask user to choose from the list.
 346       target = self.ChooseRemovableDevice(devices)
 347     else:
 348       cros_build_lib.Die('No removable devices detected.')
 349
 350     image_path = self._GetImagePath()
 351     try:
 352       self.CopyImageToDevice(image_path, self.DeviceNameToPath(target))
 353     except cros_build_lib.RunCommandError:
 354       logging.error('Failed copying image to device %s',
 355                     self.DeviceNameToPath(target))
 356
 357
 358 class FileImager(USBImager):
 359   """Copy image to the target path."""
 360
 361   def Run(self):
 362     """Copy the image to the path specified by self.device."""
 363     image_path = self._GetImagePath()
 364     if os.path.isdir(self.device):
 365       logging.info('Copying to %s',
 366                    os.path.join(self.device, os.path.basename(image_path)))
 367     else:
 368       logging.info('Copying to %s', self.device)
 369     try:
 370       shutil.copy(image_path, self.device)
 371     except IOError:
 372       logging.error('Failed to copy image %s to %s', image_path, self.device)
 373
 374
 375 class DeviceUpdateError(Exception):
 376   """Thrown when there is an error during device update."""
 377
 378
 379 class RemoteDeviceUpdater(object):
 380   """Performs update on a remote device."""
 381   ROOTFS_FILENAME = 'update.gz'
 382   STATEFUL_FILENAME = 'stateful.tgz'
 383   DEVSERVER_PKG_DIR = os.path.join(constants.SOURCE_ROOT, 'src/platform/dev')
 384   DEVSERVER_FILENAME = 'devserver.py'
 385   STATEFUL_UPDATE_BIN = '/usr/bin/stateful_update'
 386   UPDATE_ENGINE_BIN = 'update_engine_client'
 387   UPDATE_CHECK_INTERVAL = 10
 388   # Root working directory on the device. This directory is in the
 389   # stateful partition and thus has enough space to store the payloads.
 390   DEVICE_BASE_DIR = '/mnt/stateful_partition/cros-flash'
 391
 392   def __init__(self, ssh_hostname, ssh_port, image, stateful_update=True,
 393                rootfs_update=True, clobber_stateful=False, reboot=True,
 394                board=None, src_image_to_delta=None, wipe=True, debug=False,
 395                yes=False):
 396     """Initializes RemoteDeviceUpdater"""
 397     if not stateful_update and not rootfs_update:
 398       cros_build_lib.Die('No update operation to perform. Use -h to see usage.')
 399
 400     self.tempdir = tempfile.mkdtemp(prefix='cros-flash')
 401     self.ssh_hostname = ssh_hostname
 402     self.ssh_port = ssh_port
 403     self.image = image
 404     self.board = board
 405     self.src_image_to_delta = src_image_to_delta
 406     self.do_stateful_update = stateful_update
 407     self.do_rootfs_update = rootfs_update
 408     self.clobber_stateful = clobber_stateful
 409     self.reboot = reboot
 410     self.debug = debug
 411     # Do not wipe if debug is set.
 412     self.wipe = wipe and not debug
 413     self.yes = yes
 414     # The variables below are set if user passes an local image path.
 415     # Used to store a copy of the local image.
 416     self.image_tempdir = None
 417     # Used to store a symlink in devserver's static_dir.
 418     self.static_tempdir = None
 419
 420   @classmethod
 421   def GetUpdateStatus(cls, device, keys=None):
 422     """Returns the status of the update engine on the |device|.
 423
 424     Retrieves the status from update engine and confirms all keys are
 425     in the status.
 426
 427     Args:
 428       device: A ChromiumOSDevice object.
 429       keys: the keys to look for in the status result (defaults to
 430         ['CURRENT_OP']).
 431
 432     Returns:
 433       A list of values in the order of |keys|.
 434     """
 435     keys = ['CURRENT_OP'] if not keys else keys
 436     result = device.RunCommand([cls.UPDATE_ENGINE_BIN, '--status'],
 437                                capture_output=True)
 438     if not result.output:
 439       raise Exception('Cannot get update status')
 440
 441     try:
 442       status = cros_build_lib.LoadKeyValueFile(
 443           cStringIO.StringIO(result.output))
 444     except ValueError:
 445       raise ValueError('Cannot parse update status')
 446
 447     values = []
 448     for key in keys:
 449       if key not in status:
 450         raise ValueError('Missing %s in the update engine status')
 451
 452       values.append(status.get(key))
 453
 454     return values
 455
 456   def UpdateStateful(self, device, payload, clobber=False):
 457     """Update the stateful partition of the device.
 458
 459     Args:
 460       device: The ChromiumOSDevice object to update.
 461       payload: The path to the update payload.
 462       clobber: Clobber stateful partition (defaults to False).
 463     """
 464     # Copy latest stateful_update to device.
 465     stateful_update_bin = cros_build_lib.FromChrootPath(
 466         self.STATEFUL_UPDATE_BIN)
 467     device.CopyToWorkDir(stateful_update_bin)
 468     msg = 'Updating stateful partition'
 469     logging.info('Copying stateful payload to device...')
 470     device.CopyToWorkDir(payload)
 471     cmd = ['sh',
 472            os.path.join(device.work_dir,
 473                         os.path.basename(self.STATEFUL_UPDATE_BIN)),
 474            os.path.join(device.work_dir, os.path.basename(payload))]
 475
 476     if clobber:
 477       cmd.append('--stateful_change=clean')
 478       msg += ' with clobber enabled'
 479
 480     logging.info('%s...', msg)
 481     try:
 482       device.RunCommand(cmd)
 483     except cros_build_lib.RunCommandError:
 484       logging.error('Faild to perform stateful partition update.')
 485
 486   def _CopyDevServerPackage(self, device, tempdir):
 487     """Copy devserver package to work directory of device.
 488
 489     Args:
 490       device: The ChromiumOSDevice object to copy the package to.
 491       tempdir: The directory to temporarily store devserver package.
 492     """
 493     logging.info('Copying devserver package to device...')
 494     src_dir = os.path.join(tempdir, 'src')
 495     osutils.RmDir(src_dir, ignore_missing=True)
 496     shutil.copytree(
 497         self.DEVSERVER_PKG_DIR, src_dir,
 498         ignore=shutil.ignore_patterns('*.pyc', 'tmp*', '.*', 'static', '*~'))
 499     device.CopyToWorkDir(src_dir)
 500     return os.path.join(device.work_dir, os.path.basename(src_dir))
 501
 502   def SetupRootfsUpdate(self, device):
 503     """Makes sure |device| is ready for rootfs update."""
 504     logging.info('Checking if update engine is idle...')
 505     status, = self.GetUpdateStatus(device)
 506     if status == 'UPDATE_STATUS_UPDATED_NEED_REBOOT':
 507       logging.info('Device needs to reboot before updating...')
 508       device.Reboot()
 509       status, = self.GetUpdateStatus(device)
 510
 511     if status != 'UPDATE_STATUS_IDLE':
 512       raise DeviceUpdateError('Update engine is not idle. Status: %s' % status)
 513
 514   def UpdateRootfs(self, device, payload, tempdir):
 515     """Update the rootfs partition of the device.
 516
 517     Args:
 518       device: The ChromiumOSDevice object to update.
 519       payload: The path to the update payload.
 520       tempdir: The directory to store temporary files.
 521     """
 522     # Setup devserver and payload on the target device.
 523     static_dir = os.path.join(device.work_dir, 'static')
 524     payload_dir = os.path.join(static_dir, 'pregenerated')
 525     src_dir = self._CopyDevServerPackage(device, tempdir)
 526     device.RunCommand(['mkdir', '-p', payload_dir])
 527     logging.info('Copying rootfs payload to device...')
 528     device.CopyToDevice(payload, payload_dir)
 529     devserver_bin = os.path.join(src_dir, self.DEVSERVER_FILENAME)
 530     ds = ds_wrapper.RemoteDevServerWrapper(
 531         device, devserver_bin, static_dir=static_dir, log_dir=device.work_dir)
 532
 533     logging.info('Updating rootfs partition')
 534     try:
 535       ds.Start()
 536
 537       omaha_url = ds.GetDevServerURL(ip=remote_access.LOCALHOST_IP,
 538                                      port=ds.port,
 539                                      sub_dir='update/pregenerated')
 540       cmd = [self.UPDATE_ENGINE_BIN, '-check_for_update',
 541              '-omaha_url=%s' % omaha_url]
 542       device.RunCommand(cmd)
 543
 544       # Loop until update is complete.
 545       while True:
 546         op, progress = self.GetUpdateStatus(device, ['CURRENT_OP', 'PROGRESS'])
 547         logging.info('Waiting for update...status: %s at progress %s',
 548                      op, progress)
 549
 550         if op == 'UPDATE_STATUS_UPDATED_NEED_REBOOT':
 551           break
 552
 553         if op == 'UPDATE_STATUS_IDLE':
 554           raise DeviceUpdateError(
 555               'Update failed with unexpected update status: %s' % op)
 556
 557         time.sleep(self.UPDATE_CHECK_INTERVAL)
 558
 559       ds.Stop()
 560     except Exception:
 561       logging.error('Rootfs update failed.')
 562       logging.warning(ds.TailLog() or 'No devserver log is available.')
 563       raise
 564     finally:
 565       ds.Stop()
 566       device.CopyFromDevice(ds.log_filename,
 567                             os.path.join(tempdir, 'target_devserver.log'),
 568                             error_code_ok=True)
 569       device.CopyFromDevice('/var/log/update_engine.log', tempdir,
 570                             follow_symlinks=True,
 571                             error_code_ok=True)
 572
 573   def ConvertLocalPathToXbuddyPath(self, path):
 574     """Converts |path| to an xbuddy path.
 575
 576     This function copies the image into a temprary directory in chroot
 577     and creates a symlink in static_dir for devserver/xbuddy to
 578     access.
 579
 580     Args:
 581       path: Path to an image.
 582
 583     Returns:
 584       The xbuddy path for |path|.
 585     """
 586     self.image_tempdir = osutils.TempDir(
 587         base_dir=cros_build_lib.FromChrootPath('/tmp'),
 588         prefix='cros_flash_local_image',
 589         sudo_rm=True)
 590
 591     tempdir_path = self.image_tempdir.tempdir
 592     logging.info('Copying image to temporary directory %s', tempdir_path)
 593     # Devserver only knows the image names listed in IMAGE_TYPE_TO_NAME.
 594     # Rename the image to chromiumos_test_image.bin when copying.
 595     TEMP_IMAGE_TYPE  = 'test'
 596     shutil.copy(path,
 597                 os.path.join(tempdir_path, IMAGE_TYPE_TO_NAME[TEMP_IMAGE_TYPE]))
 598     chroot_path = cros_build_lib.ToChrootPath(tempdir_path)
 599     # Create and link static_dir/local_imagexxxx/link to the image
 600     # folder, so that xbuddy/devserver can understand the path.
 601     # Alternatively, we can to pass '--image' at devserver startup,
 602     # but this flag is deprecated.
 603     self.static_tempdir = osutils.TempDir(base_dir=DEVSERVER_STATIC_DIR,
 604                                           prefix='local_image',
 605                                           sudo_rm=True)
 606     relative_dir = os.path.join(os.path.basename(self.static_tempdir.tempdir),
 607                                 'link')
 608     symlink_path = os.path.join(DEVSERVER_STATIC_DIR, relative_dir)
 609     logging.info('Creating a symlink %s -> %s', symlink_path, chroot_path)
 610     os.symlink(chroot_path, symlink_path)
 611     return os.path.join(relative_dir, TEMP_IMAGE_TYPE)
 612
 613   def GetUpdatePayloads(self, path, payload_dir, board=None,
 614                         src_image_to_delta=None, timeout=60 * 15):
 615     """Launch devserver to get the update payloads.
 616
 617     Args:
 618       path: The xbuddy path.
 619       payload_dir: The directory to store the payloads.
 620       board: The default board to use when |path| is None.
 621       src_image_to_delta: Image used as the base to generate the delta payloads.
 622       timeout: Timeout for launching devserver (seconds).
 623     """
 624     ds = ds_wrapper.DevServerWrapper(static_dir=DEVSERVER_STATIC_DIR,
 625                                      src_image=src_image_to_delta, board=board)
 626     req = GenerateXbuddyRequest(path, 'update')
 627     logging.info('Starting local devserver to generate/serve payloads...')
 628     try:
 629       ds.Start()
 630       url = ds.OpenURL(ds.GetDevServerURL(sub_dir=req), timeout=timeout)
 631       ds.DownloadFile(os.path.join(url, self.ROOTFS_FILENAME), payload_dir)
 632       ds.DownloadFile(os.path.join(url, self.STATEFUL_FILENAME), payload_dir)
 633     except ds_wrapper.DevServerException:
 634       logging.warning(ds.TailLog() or 'No devserver log is available.')
 635       raise
 636     else:
 637       logging.debug(ds.TailLog() or 'No devserver log is available.')
 638     finally:
 639       ds.Stop()
 640       if os.path.exists(ds.log_filename):
 641         shutil.copyfile(ds.log_filename,
 642                         os.path.join(payload_dir, 'local_devserver.log'))
 643       else:
 644         logging.warning('Could not find %s', ds.log_filename)
 645
 646   def _CheckPayloads(self, payload_dir):
 647     """Checks that all update payloads exists in |payload_dir|."""
 648     filenames = []
 649     filenames += [self.ROOTFS_FILENAME] if self.do_rootfs_update else []
 650     filenames += [self.STATEFUL_FILENAME] if self.do_stateful_update else []
 651     for fname in filenames:
 652       payload = os.path.join(payload_dir, fname)
 653       if not os.path.exists(payload):
 654         cros_build_lib.Die('Payload %s does not exist!' % payload)
 655
 656   def Verify(self, old_root_dev, new_root_dev):
 657     """Verifies that the root deivce changed after reboot."""
 658     assert new_root_dev and old_root_dev
 659     if new_root_dev == old_root_dev:
 660       raise DeviceUpdateError(
 661           'Failed to boot into the new version. Possibly there was a '
 662           'signing problem, or an automated rollback occurred because '
 663           'your new image failed to boot.')
 664
 665   @classmethod
 666   def GetRootDev(cls, device):
 667     """Get the current root device on |device|."""
 668     rootdev = device.RunCommand(
 669         ['rootdev', '-s'], capture_output=True).output.strip()
 670     logging.debug('Current root device is %s', rootdev)
 671     return rootdev
 672
 673   def Cleanup(self):
 674     """Cleans up the temporary directory."""
 675     if self.image_tempdir:
 676       self.image_tempdir.Cleanup()
 677
 678     if self.static_tempdir:
 679       self.static_tempdir.Cleanup()
 680
 681     if self.wipe:
 682       logging.info('Cleaning up temporary working directory...')
 683       osutils.RmDir(self.tempdir)
 684     else:
 685       logging.info('You can find the log files and/or payloads in %s',
 686                    self.tempdir)
 687
 688   def _CanRunDevserver(self, device, tempdir):
 689     """We can run devserver on |device|.
 690
 691     If the stateful partition is corrupted, Python or other packages
 692     (e.g. cherrypy) that Cros Flash needs for rootfs update may be
 693     missing on |device|.
 694
 695     Args:
 696        device: A ChromiumOSDevice object.
 697        tempdir: A temporary directory to store files.
 698
 699     Returns:
 700       True if we can start devserver; False otherwise.
 701     """
 702     logging.info('Checking if we can run devserver on the device.')
 703     src_dir = self._CopyDevServerPackage(device, tempdir)
 704     devserver_bin = os.path.join(src_dir, self.DEVSERVER_FILENAME)
 705     try:
 706       device.RunCommand(['python', devserver_bin, '--help'])
 707     except cros_build_lib.RunCommandError as e:
 708       logging.warning('Cannot start devserver: %s', e)
 709       return False
 710
 711     return True
 712
 713   def Run(self):
 714     """Performs remote device update."""
 715     old_root_dev, new_root_dev = None, None
 716     try:
 717       with remote_access.ChromiumOSDeviceHandler(
 718           self.ssh_hostname, port=self.ssh_port,
 719           base_dir=self.DEVICE_BASE_DIR) as device:
 720
 721         board = cros_build_lib.GetBoard(device_board=device.board,
 722                                         override_board=self.board,
 723                                         force=self.yes)
 724         logging.info('Board is %s', board)
 725
 726         if os.path.isdir(self.image):
 727           # If the given path is a directory, we use the provided
 728           # update payload(s) in the directory.
 729           payload_dir = self.image
 730           logging.info('Using provided payloads in %s', payload_dir)
 731         else:
 732           if os.path.isfile(self.image):
 733             # If the given path is an image, make sure devserver can
 734             # access it and generate payloads.
 735             logging.info('Using image %s', self.image)
 736             image_path = self.ConvertLocalPathToXbuddyPath(self.image)
 737           else:
 738             # For xbuddy paths, we should do a sanity check / confirmation
 739             # when the xbuddy board doesn't match the board on the
 740             # device. Unfortunately this isn't currently possible since we
 741             # don't want to duplicate xbuddy code.  TODO(sosa):
 742             # crbug.com/340722 and use it to compare boards.
 743
 744             # Translate the xbuddy path to get the exact image to use.
 745             translated_path = TranslateImagePath(self.image, board)
 746             logging.info('Using image %s', translated_path)
 747             # Convert the translated path to be used in the update request.
 748             image_path = ConvertTranslatedPath(self.image, translated_path)
 749
 750           # Launch a local devserver to generate/serve update payloads.
 751           payload_dir = self.tempdir
 752           self.GetUpdatePayloads(image_path, payload_dir,
 753                                  board=board,
 754                                  src_image_to_delta=self.src_image_to_delta)
 755
 756         # Verify that all required payloads are in the payload directory.
 757         self._CheckPayloads(payload_dir)
 758
 759         restore_stateful = False
 760         if (not self._CanRunDevserver(device, self.tempdir) and
 761             self.do_rootfs_update):
 762           msg = ('Cannot start devserver! The stateful partition may be '
 763                  'corrupted. Cros Flash can try to restore the stateful '
 764                  'partition first.')
 765           restore_stateful = self.yes or cros_build_lib.BooleanPrompt(
 766               default=False, prolog=msg)
 767           if not restore_stateful:
 768             cros_build_lib.Die('Cannot continue to perform rootfs update!')
 769
 770         if restore_stateful:
 771           logging.warning('Restoring the stateful partition...')
 772           payload = os.path.join(payload_dir, self.STATEFUL_FILENAME)
 773           self.UpdateStateful(device, payload, clobber=self.clobber_stateful)
 774           device.Reboot()
 775           if self._CanRunDevserver(device, self.tempdir):
 776             logging.info('Stateful partition restored.')
 777           else:
 778             cros_build_lib.Die('Unable to restore stateful partition. Exiting.')
 779
 780         # Perform device updates.
 781         if self.do_rootfs_update:
 782           self.SetupRootfsUpdate(device)
 783           # Record the current root device. This must be done after
 784           # SetupRootfsUpdate because SetupRootfsUpdate may reboot the
 785           # device if there is a pending update, which changes the
 786           # root device.
 787           old_root_dev = self.GetRootDev(device)
 788           payload = os.path.join(payload_dir, self.ROOTFS_FILENAME)
 789           self.UpdateRootfs(device, payload, self.tempdir)
 790           logging.info('Rootfs update completed.')
 791
 792         if self.do_stateful_update and not restore_stateful:
 793           payload = os.path.join(payload_dir, self.STATEFUL_FILENAME)
 794           self.UpdateStateful(device, payload, clobber=self.clobber_stateful)
 795           logging.info('Stateful update completed.')
 796
 797         if self.reboot:
 798           logging.info('Rebooting device..')
 799           device.Reboot()
 800           if self.clobber_stateful:
 801             # --clobber-stateful wipes the stateful partition and the
 802             # working directory on the device no longer exists. To
 803             # remedy this, we recreate the working directory here.
 804             device.BaseRunCommand(['mkdir', '-p', device.work_dir])
 805
 806         if self.do_rootfs_update and self.reboot:
 807           new_root_dev = self.GetRootDev(device)
 808           self.Verify(old_root_dev, new_root_dev)
 809
 810     except Exception:
 811       logging.error('Device update failed.')
 812       raise
 813     else:
 814       logging.info('Update performed successfully.')
 815     finally:
 816       self.Cleanup()
 817
 818
 819 @cros.CommandDecorator('flash')
 820 class FlashCommand(cros.CrosCommand):
 821   """Update the device with an image.
 822
 823   This command updates the device with the image
 824   (ssh://<hostname>:{port}, copies an image to a removable device
 825   (usb://<device_path), or copies a xbuddy path to a local
 826   file path with (file://file_path).
 827
 828   For device update, it assumes that device is able to accept ssh
 829   connections.
 830
 831   For rootfs partition update, this command may launch a devserver to
 832   generate payloads. As a side effect, it may create symlinks in
 833   static_dir/others used by the devserver.
 834   """
 835
 836   EPILOG = """
 837 To update/image the device with the latest locally built image:
 838   cros flash device latest
 839   cros flash device
 840
 841 To update/image the device with an xbuddy path:
 842   cros flash device xbuddy://{local, remote}/<board>/<version>
 843
 844   Common xbuddy version aliases are 'latest' (alias for 'latest-stable')
 845   latest-{dev, beta, stable, canary}, and latest-official.
 846
 847 To update/image the device with a local image path:
 848   cros flash device /path/to/image.bin
 849
 850 Examples:
 851   cros flash 192.168.1.7 xbuddy://remote/x86-mario/latest-canary
 852   cros flash 192.168.1.7 xbuddy://remote/x86-mario-paladin/R32-4830.0.0-rc1
 853   cros flash usb:// xbuddy://remote/trybot-x86-mario-paladin/R32-5189.0.0-b100
 854
 855   For more information and known problems/fixes, please see:
 856   http://dev.chromium.org/chromium-os/build/cros-flash
 857 """
 858
 859   SSH_MODE = 'ssh'
 860   USB_MODE = 'usb'
 861   FILE_MODE = 'file'
 862
 863   # Override base class property to enable stats upload.
 864   upload_stats = True
 865
 866   @classmethod
 867   def AddParser(cls, parser):
 868     """Add parser arguments."""
 869     super(FlashCommand, cls).AddParser(parser)
 870     parser.add_argument(
 871         'device', help='ssh://device_hostname[:port] or usb://{device_path}. '
 872         'If no device_path is given (i.e. usb://), user will be prompted to '
 873         'choose from a list of removable devices.')
 874     parser.add_argument(
 875         'image', nargs='?', default='latest', help="A local path or an xbuddy "
 876         "path: xbuddy://{local|remote}/board/version/{image_type} image_type "
 877         "can be: 'test', 'dev', 'base', or 'recovery'. Note any strings that "
 878         "do not map to a real file path will be converted to an xbuddy path "
 879         "i.e., latest, will map to xbuddy://latest.")
 880     parser.add_argument(
 881         '--clear-cache', default=False, action='store_true',
 882         help='Clear the devserver static directory. This deletes all the '
 883         'downloaded images and payloads, and also payloads generated by '
 884         'the devserver. Default is not to clear.')
 885
 886     update = parser.add_argument_group('Advanced device update options')
 887     update.add_argument(
 888         '--board', default=None, help='The board to use. By default it is '
 889         'automatically detected. You can override the detected board with '
 890         'this option')
 891     update.add_argument(
 892         '--yes', default=False, action='store_true',
 893         help='Force yes to any prompt. Use with caution.')
 894     update.add_argument(
 895         '--no-reboot', action='store_false', dest='reboot', default=True,
 896         help='Do not reboot after update. Default is always reboot.')
 897     update.add_argument(
 898         '--no-wipe', action='store_false', dest='wipe', default=True,
 899         help='Do not wipe the temporary working directory. Default '
 900         'is always wipe.')
 901     update.add_argument(
 902         '--no-stateful-update', action='store_false', dest='stateful_update',
 903         help='Do not update the stateful partition on the device. '
 904         'Default is always update.')
 905     update.add_argument(
 906         '--no-rootfs-update', action='store_false', dest='rootfs_update',
 907         help='Do not update the rootfs partition on the device. '
 908         'Default is always update.')
 909     update.add_argument(
 910         '--src-image-to-delta', type='path',
 911         help='Local path to an image to be used as the base to generate '
 912         'delta payloads.')
 913     update.add_argument(
 914         '--clobber-stateful', action='store_true', default=False,
 915         help='Clobber stateful partition when performing update.')
 916
 917   def __init__(self, options):
 918     """Initializes cros flash."""
 919     cros.CrosCommand.__init__(self, options)
 920     self.run_mode = None
 921     self.ssh_hostname = None
 922     self.ssh_port = None
 923     self.usb_dev = None
 924     self.copy_path = None
 925     self.any = False
 926
 927   def _ParseDevice(self, device):
 928     """Parse |device| and set corresponding variables ."""
 929     # pylint: disable=E1101
 930     if urlparse.urlparse(device).scheme == '':
 931       # For backward compatibility, prepend ssh:// ourselves.
 932       device = 'ssh://%s' % device
 933
 934     parsed = urlparse.urlparse(device)
 935     if parsed.scheme == self.SSH_MODE:
 936       self.run_mode = self.SSH_MODE
 937       self.ssh_hostname = parsed.hostname
 938       self.ssh_port = parsed.port
 939     elif parsed.scheme == self.USB_MODE:
 940       self.run_mode = self.USB_MODE
 941       self.usb_dev = device[len('%s://' % self.USB_MODE):]
 942     elif parsed.scheme == self.FILE_MODE:
 943       self.run_mode = self.FILE_MODE
 944       self.copy_path = device[len('%s://' % self.FILE_MODE):]
 945     else:
 946       cros_build_lib.Die('Does not support device %s' % device)
 947
 948   # pylint: disable=E1101
 949   def Run(self):
 950     """Perfrom the cros flash command."""
 951     self.options.Freeze()
 952
 953     if self.options.clear_cache:
 954       logging.info('Clearing the cache...')
 955       ds_wrapper.DevServerWrapper.WipeStaticDirectory(DEVSERVER_STATIC_DIR)
 956
 957     try:
 958       osutils.SafeMakedirsNonRoot(DEVSERVER_STATIC_DIR)
 959     except OSError:
 960       logging.error('Failed to create %s', DEVSERVER_STATIC_DIR)
 961
 962     self._ParseDevice(self.options.device)
 963     try:
 964       if self.run_mode == self.SSH_MODE:
 965         logging.info('Preparing to update the remote device %s',
 966                      self.options.device)
 967         updater = RemoteDeviceUpdater(
 968             self.ssh_hostname,
 969             self.ssh_port,
 970             self.options.image,
 971             board=self.options.board,
 972             src_image_to_delta=self.options.src_image_to_delta,
 973             rootfs_update=self.options.rootfs_update,
 974             stateful_update=self.options.stateful_update,
 975             clobber_stateful=self.options.clobber_stateful,
 976             reboot=self.options.reboot,
 977             wipe=self.options.wipe,
 978             debug=self.options.debug,
 979             yes=self.options.yes)
 980
 981         # Perform device update.
 982         updater.Run()
 983       elif self.run_mode == self.USB_MODE:
 984         logging.info('Preparing to image the removable device %s',
 985                      self.options.device)
 986         imager = USBImager(self.usb_dev,
 987                            self.options.board,
 988                            self.options.image,
 989                            debug=self.options.debug,
 990                            yes=self.options.yes)
 991         imager.Run()
 992       elif self.run_mode == self.FILE_MODE:
 993         logging.info('Preparing to copy image to %s', self.copy_path)
 994         imager = FileImager(self.copy_path,
 995                             self.options.board,
 996                             self.options.image,
 997                             debug=self.options.debug,
 998                             yes=self.options.yes)
 999         imager.Run()
1000
1001     except (Exception, KeyboardInterrupt) as e:
1002       logging.error(e)
1003       logging.error('Cros Flash failed before completing.')
1004       if self.options.debug:
1005         raise
1006     else:
1007       logging.info('Cros Flash completed successfully.')