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