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