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