bisect-kit: integrate quick-provision

If you need to specify ssh keys of devserver for running
quick-provision, the keys could be stored in the
./cros_template_files/at_home/.ssh directory and bisect-kit will copy
them into chroot.

BUG=b:170635511
TEST=manually test some of following combinations
TEST=  {switch_cros_prebuilt.py, switch_cros_localbuild.py, switch_cros_localbuild_buildbucket.py}
TEST=  {quick-provision, old cros flash, new cros flash}
TEST=  {DUT in the lab, local DUT in corp network, local DUT in public network}
TEST=  {developer workstation, bisector runner}

Change-Id: Ic2fdea88bd62b7dca2a59b27aada518f1b2cdc70
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/bisect-kit/+/2718213
Commit-Queue: Kuang-che Wu <kcwu@chromium.org>
Tested-by: Kuang-che Wu <kcwu@chromium.org>
Reviewed-by: Zheng-Jie Chang <zjchang@chromium.org>
diff --git a/bisect_kit/cros_quick_provision.py b/bisect_kit/cros_quick_provision.py
deleted file mode 100644
index 75b444b..0000000
--- a/bisect_kit/cros_quick_provision.py
+++ /dev/null
@@ -1,61 +0,0 @@
-# -*- coding: utf-8 -*-
-# Copyright 2021 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-"""ChromeOS quick provision helper.
-
-Use to provision a DUT using quick-provision script safely.
-
-TODO(kimjae): Transition to using TLS ProvisionDut for F20.
-"""
-
-import logging
-import os
-import re
-import socket
-
-from bisect_kit import cros_lab_util
-from bisect_kit import cros_util
-from bisect_kit import errors
-from bisect_kit import util
-
-logger = logging.getLogger(__name__)
-
-
-_QUICK_PROVISION_SCRIPT = 'quick-provision'
-re_chromeos_build= r'^[a-zA-Z-]+\/R\d+-[\.\d-]+$'
-
-
-def argtype_cros_build(s):
-  if not bool(re.match(re_chromeos_build, s)):
-    msg = 'invalid cros build'
-    raise cli.ArgTypeError(
-        msg, 'eve-release/R90-13766.0.0, eve-cq/R90-13766.0.0-XXX-XXX, etc')
-  return s
-
-
-def quick_provision(chromeos_root, host, build):
-  """Runs the quick-provision script on the device and
-  returns the device in a good state in case of errors.
-
-  Args:
-    chromeos_root: chromeos tree root
-    host: DUT name/address
-    build: build to provision to
-
-  Raises:
-    errors.ExternalError: quick-provision fails
-  """
-  logger.info("Starting quick-provision")
-  autotest_path = os.path.join(cros_util.chromeos_root_inside_chroot,
-                               'src/third_party/autotest/files')
-  quick_provision_cmd = ['test_that',
-                         '--args="value=\'%s\'"' % build,
-                         host,
-                         'provision_QuickProvision',
-                         '--autotest_dir',
-                         autotest_path]
-  try:
-    cros_util.cros_sdk(chromeos_root, *quick_provision_cmd)
-  except subprocess.CalledProcessError as e:
-    raise errors.ExternalError('quick-provision failed') from e
diff --git a/bisect_kit/cros_util.py b/bisect_kit/cros_util.py
index 5b338b3..b45fbfc 100644
--- a/bisect_kit/cros_util.py
+++ b/bisect_kit/cros_util.py
@@ -16,6 +16,7 @@
 import ast
 import calendar
 import datetime
+import enum
 import errno
 import glob
 import json
@@ -25,13 +26,16 @@
 import shutil
 import subprocess
 import sys
+import tempfile
 import time
+import urllib.parse
 
 from google.protobuf import json_format
 
 from bisect_kit import buildbucket_util
 from bisect_kit import cli
 from bisect_kit import codechange
+from bisect_kit import common
 from bisect_kit import cr_util
 from bisect_kit import errors
 from bisect_kit import git_util
@@ -45,6 +49,7 @@
 re_chromeos_short_version = r'^\d+\.\d+\.\d+$'
 re_chromeos_snapshot_version = r'^R\d+-\d+\.\d+\.\d+-\d+$'
 
+gs_archive_base = 'gs://chromeos-image-archive/'
 gs_archive_path = 'gs://chromeos-image-archive/{board}-release'
 
 # Assume gsutil is in PATH.
@@ -85,11 +90,14 @@
 
 chromeos_root_inside_chroot = '/mnt/host/source'
 # relative to chromeos_root
+in_tree_autotest_dir = 'src/third_party/autotest/files'
 prebuilt_autotest_dir = 'tmp/autotest-prebuilt'
 prebuilt_tast_dir = 'tmp/tast-prebuilt'
-# Relative to chromeos root. Images are cached_images_dir/$board/$image_name.
-cached_images_dir = 'src/build/images'
+# Relative to chromeos root. Images are build_images_dir/$board/$image_name.
+build_images_dir = 'src/build/images'
+cached_images_dir = 'tmp/images'
 test_image_filename = 'chromiumos_test_image.bin'
+sample_partition_filename = 'full_dev_part_KERN.bin.gz'
 
 VERSION_KEY_CROS_SHORT_VERSION = 'cros_short_version'
 VERSION_KEY_CROS_FULL_VERSION = 'cros_full_version'
@@ -99,6 +107,41 @@
 VERSION_KEY_ANDROID_BRANCH = 'android_branch'
 CROSLAND_URL_TEMPLATE = 'https://crosland.corp.google.com/log/%s..%s'
 
+autotest_shadow_config = """
+[CROS]
+enable_ssh_tunnel_for_servo: True
+enable_ssh_tunnel_for_chameleon: True
+enable_ssh_connection_for_devserver: True
+enable_ssh_tunnel_for_moblab: True
+"""
+
+
+class ImageType(enum.Enum):
+  """Chrome OS image type
+
+  It describes the image format, not image location.
+  """
+  # Full disk image like chromiumos_test_image.bin and
+  # chromiumos_test_image.tar.xz.
+  # Supported by 'cros flash'.
+  DISK_IMAGE = enum.auto()
+  # Contains files like full_dev_part_KERN.bin.gz.
+  # Supported by quick-provision and newer 'cros flash'.
+  PARTITION_IMAGE = enum.auto()
+  # Contains files like image.zip. We need to unzip first.
+  ZIP_FILE = enum.auto()
+
+
+class ImageInfo(dict):
+  """Image info (dict: image type -> path).
+
+  For a given Chrome OS version, there are several image formats available.
+  This class describes a collection of images for a certain Chrome OS version.
+  cros_flash() or quick_provision() can resolve a compatible image from this
+  object. `image type` is an ImageType enum. `path` could be a path on the
+  local disk or a remote URI.
+  """
+
 
 class NeedRecreateChrootException(Exception):
   """Failed to build ChromeOS because of chroot mismatch or corruption"""
@@ -422,22 +465,6 @@
   raise errors.ExternalError('reboot failed?')
 
 
-def gs_release_boardpath(board):
-  """Normalizes board name for gs://chromeos-releases/
-
-  This follows behavior of PushImage() in chromite/scripts/pushimage.py
-  Note, only gs://chromeos-releases/ needs normalization,
-  gs://chromeos-image-archive does not.
-
-  Args:
-    board: ChromeOS board name
-
-  Returns:
-    normalized board name
-  """
-  return board.replace('_', '-')
-
-
 def gsutil(*args, **kwargs):
   """gsutil command line wrapper.
 
@@ -852,54 +879,181 @@
   return result
 
 
-def prepare_snapshot_image(chromeos_root, board, snapshot_version):
-  """Prepare chromeos snapshot image.
+def search_snapshot_image(board, snapshot_version):
+  """Searches chromeos snapshot image.
 
   Args:
-    chromeos_root: chromeos tree root
     board: ChromeOS board name
     snapshot_version: ChromeOS snapshot version number
 
   Returns:
-    local file path of test image relative to chromeos_root
+    ImageInfo object
   """
   assert is_cros_snapshot_version(snapshot_version)
-  milestone, short_version, snapshot_id = snapshot_version_split(
-      snapshot_version)
-  full_version = make_cros_full_version(milestone, short_version)
-  tmp_dir = os.path.join(
-      chromeos_root, 'tmp',
-      'ChromeOS-test-%s-%s-%s' % (full_version, board, snapshot_id))
-  if not os.path.exists(tmp_dir):
-    os.makedirs(tmp_dir)
+  image_info = ImageInfo()
+  gs_path = gs_archive_base + '{board}-snapshot/{snapshot_version}-*'.format(
+      board=board, snapshot_version=snapshot_version)
+  files = gsutil_ls(
+      gs_path + '/' + sample_partition_filename, ignore_errors=True)
+  if files:
+    image_info[ImageType.PARTITION_IMAGE] = files[0].replace(
+        sample_partition_filename, '')
+  files = gsutil_ls(gs_path + '/image.zip', ignore_errors=True)
+  if files:
+    image_info[ImageType.ZIP_FILE] = files[0]
+  return image_info
 
-  gs_path = ('gs://chromeos-image-archive/{board}-snapshot/' +
-             '{snapshot_version}-*/image.zip')
-  gs_path = gs_path.format(board=board, snapshot_version=snapshot_version)
 
-  full_path = os.path.join(tmp_dir, test_image_filename)
-  rel_path = os.path.relpath(full_path, chromeos_root)
-  if os.path.exists(full_path):
-    return rel_path
+def prepare_image_for_quick_provision(image_info):
+  path = image_info.get(ImageType.PARTITION_IMAGE)
+  if path and path.startswith(gs_archive_base):
+    return urllib.parse.urlparse(path).path[1:]
 
-  files = gsutil_ls(gs_path, ignore_errors=True)
-  if len(files) >= 1:
-    gs_path = files[0]
-    gsutil('cp', gs_path, tmp_dir)
-    util.check_call(
-        'unzip', '-j', 'image.zip', test_image_filename, cwd=tmp_dir)
-    os.remove(os.path.join(tmp_dir, 'image.zip'))
-    return rel_path
-
-  assert False
+  logger.warning(
+      'image format or location are not supported by quick-provision: %s',
+      image_info)
   return None
 
 
-def prepare_prebuilt_image(chromeos_root, board, version):
-  """Prepare chromeos prebuilt image.
+def _cache_path_for_download(chromeos_root, url):
+  cache_folder = os.path.join(chromeos_root, cached_images_dir)
+  os.makedirs(cache_folder, exist_ok=True)
+  name = urllib.parse.quote(url, safe='')
+  return os.path.join(cache_folder, name)
 
-  It searches for xbuddy image which "cros flash" can use, or fetch image to
-  local disk.
+
+def prepare_image_for_cros_flash(chromeos_root, image_info):
+  path = image_info.get(ImageType.DISK_IMAGE)
+  if path:
+    # local path
+    if '://' not in path:
+      return path
+
+    m = re.search(gs_archive_base + r'([^/]+)-[a-z]+/([^/]+)/', path)
+    if m:
+      return 'xbuddy://remote/%s/%s/test' % (m.group(1), m.group(2))
+    if path.startswith(gs_archive_base):
+      return path.replace('chromiumos_test_image.tar.xz', 'test')
+
+    # 'cros flash' doesn't support other gs bucket, download to local.
+    if path.startswith('gs://'):
+      cache_path = _cache_path_for_download(chromeos_root, path)
+      if os.path.exists(cache_path):
+        return cache_path
+      gsutil('cp', path, cache_path)
+      return cache_path
+
+  path = image_info.get(ImageType.PARTITION_IMAGE)
+  if path and path.startswith(gs_archive_base):
+    # newer 'cros flash' support partition images
+    # TODO(kcwu): update the commit hash after new flash is enabled
+    if git_util.is_ancestor_commit(
+        os.path.join(chromeos_root, 'chromite'), '153f91684a9a59', 'HEAD'):
+      return path
+
+  path = image_info.get(ImageType.ZIP_FILE)
+  if path:
+    cached_image_full_path = _cache_path_for_download(
+        chromeos_root, path + '.' + test_image_filename)
+    if os.path.exists(cached_image_full_path):
+      return cached_image_full_path
+
+    tmp_dir = tempfile.mkdtemp()
+    try:
+      if path.startswith('gs://'):
+        gsutil('cp', path, tmp_dir)
+        path = os.path.join(tmp_dir, os.path.basename(path))
+        assert os.path.exists(path)
+      util.check_call('unzip', '-j', path, test_image_filename, cwd=tmp_dir)
+      shutil.move(
+          os.path.join(tmp_dir, test_image_filename), cached_image_full_path)
+    finally:
+      shutil.rmtree(tmp_dir)
+
+    return cached_image_full_path
+
+  return None
+
+
+def quick_provision(chromeos_root, host, image_info):
+  # TODO(kimjae): Transition to using TLS ProvisionDut for F20.
+  logger.debug('quick_provision %s %s', host, image_info)
+  build = prepare_image_for_quick_provision(image_info)
+  if not build:
+    return False
+
+  autotest_path = os.path.join(chromeos_root_inside_chroot,
+                               in_tree_autotest_dir)
+  quick_provision_cmd = [
+      'test_that', '--args',
+      "value='%s'" % build, host, 'provision_QuickProvision', '--autotest_dir',
+      autotest_path
+  ]
+  try:
+    cros_sdk(chromeos_root, *quick_provision_cmd)
+  except subprocess.CalledProcessError as e:
+    raise errors.ExternalError('quick-provision failed') from e
+  return True
+
+
+def verify_dut_version(host, board, version):
+  if version:
+    # In the past, cros flash may fail with returncode=0
+    # So let's have an extra check.
+    if is_cros_snapshot_version(version):
+      builder_path = query_dut_lsb_release(host).get(
+          'CHROMEOS_RELEASE_BUILDER_PATH', '')
+      expect_prefix = '%s-snapshot/%s-' % (board, version)
+      if not builder_path.startswith(expect_prefix):
+        raise errors.ExternalError(
+            'although provision succeeded, the OS builder path is '
+            'unexpected: actual=%s expect=%s' % (builder_path, expect_prefix))
+    else:
+      expect_version = version_to_short(version)
+      dut_version = query_dut_short_version(host)
+      if dut_version != expect_version:
+        raise errors.ExternalError(
+            'although provision succeeded, the OS version is unexpected: '
+            'actual=%s expect=%s' % (dut_version, expect_version))
+
+  # "cros flash" may terminate successfully but the DUT starts self-repairing
+  # (b/130786578), so it's necessary to do sanity check.
+  if not is_good_dut(host):
+    raise errors.ExternalError(
+        'although provision succeeded, the DUT is in bad state')
+
+
+def provision_image(chromeos_root,
+                    host,
+                    board,
+                    image_info,
+                    version=None,
+                    clobber_stateful=False,
+                    disable_rootfs_verification=True,
+                    force_reboot_callback=None):
+  # Try quick_provision first, but fallback to cros flash.
+  # TODO(kcwu): only use quick_provision for DUTs in the lab
+  try:
+    if quick_provision(chromeos_root, host, image_info):
+      verify_dut_version(host, board, version)
+      return
+    logger.debug('quick-provision is not supported; fallback to cros flash')
+  except errors.ExternalError as e:
+    logger.warning('quick-provision failed; fallback to cros flash: %s', e)
+
+  if not cros_flash(
+      chromeos_root,
+      host,
+      image_info,
+      clobber_stateful=clobber_stateful,
+      disable_rootfs_verification=disable_rootfs_verification,
+      force_reboot_callback=force_reboot_callback):
+    raise errors.InternalError('unsupported image: ' + str(image_info))
+  verify_dut_version(host, board, version)
+
+
+def search_prebuilt_image(board, version):
+  """Searches chromeos prebuilt image.
 
   Args:
     chromeos_root: chromeos tree root
@@ -907,24 +1061,29 @@
     version: ChromeOS version number in short or full format
 
   Returns:
-    xbuddy path or None
+    ImageInfo object
   """
-  del chromeos_root  # unused
   assert is_cros_version(version)
   full_version = version_to_full(board, version)
 
+  image_info = ImageInfo()
   gs_path = gs_archive_path.format(board=board) + '/' + full_version
-  if gsutil_ls('-d', gs_path, ignore_errors=True):
-    return 'xbuddy://remote/{board}/{full_version}/test'.format(
-        board=board, full_version=full_version)
-  return None
+  if gsutil_ls(gs_path + '/chromiumos_test_image.tar.xz', ignore_errors=True):
+    image_info[ImageType.DISK_IMAGE] = gs_path + '/chromiumos_test_image.tar.xz'
+  if gsutil_ls(gs_path + '/' + sample_partition_filename, ignore_errors=True):
+    image_info[ImageType.PARTITION_IMAGE] = gs_path
+  return image_info
+
+
+def search_image(board, version):
+  if is_cros_snapshot_version(version):
+    return search_snapshot_image(board, version)
+  return search_prebuilt_image(board, version)
 
 
 def cros_flash(chromeos_root,
                host,
-               board,
-               image_path,
-               version=None,
+               image_info,
                clobber_stateful=False,
                disable_rootfs_verification=True,
                force_reboot_callback=None):
@@ -936,18 +1095,20 @@
     chromeos_root: use 'cros flash' of which chromeos tree
     host: DUT address
     board: ChromeOS board name
-    image_path: chromeos image xbuddy path or file path. For relative
-        path, it should be relative to chromeos_root.
+    image_info: ImageInfo object
     version: ChromeOS version in short or full format
     clobber_stateful: Clobber stateful partition when performing update
     disable_rootfs_verification: Disable rootfs verification after update
         is completed
     force_reboot_callback: powerful reboot hook (via servo)
 
+  Returns:
+    False for unsupported images
+
   Raises:
     errors.ExternalError: cros flash failed
   """
-  logger.info('cros_flash %s %s %s %s', host, board, version, image_path)
+  logger.info('cros_flash %s %s', host, image_info)
 
   # Reboot is necessary because sometimes previous 'cros flash' failed and
   # entered a bad state.
@@ -967,6 +1128,10 @@
     except subprocess.CalledProcessError:
       pass  # not started; do nothing
 
+  image_path = prepare_image_for_cros_flash(chromeos_root, image_info)
+  if not image_path:
+    return False
+
   # Handle relative path.
   if '://' not in image_path and not os.path.isabs(image_path):
     assert os.path.exists(os.path.join(chromeos_root, image_path))
@@ -991,46 +1156,27 @@
   if disable_rootfs_verification:
     args.append('--disable-rootfs-verification')
 
+  # TODO(kcwu): remove this flag once the new flash code enabled by default
+  if git_util.is_ancestor_commit(
+      os.path.join(chromeos_root, 'chromite'), '153f91684a9a59', 'HEAD'):
+    args.append('--exp-new-flash')
+
   try:
     cros_sdk(chromeos_root, 'cros', 'flash', *args)
   except subprocess.CalledProcessError as e:
     raise errors.ExternalError('cros flash failed') from e
-
-  if version:
-    # In the past, cros flash may fail with returncode=0
-    # So let's have an extra check.
-    if is_cros_snapshot_version(version):
-      builder_path = query_dut_lsb_release(host).get(
-          'CHROMEOS_RELEASE_BUILDER_PATH', '')
-      expect_prefix = '%s-snapshot/%s-' % (board, version)
-      if not builder_path.startswith(expect_prefix):
-        raise errors.ExternalError(
-            'although cros flash succeeded, the OS builder path is '
-            'unexpected: actual=%s expect=%s' % (builder_path, expect_prefix))
-    else:
-      expect_version = version_to_short(version)
-      dut_version = query_dut_short_version(host)
-      if dut_version != expect_version:
-        raise errors.ExternalError(
-            'although cros flash succeeded, the OS version is unexpected: '
-            'actual=%s expect=%s' % (dut_version, expect_version))
-
-  # "cros flash" may terminate successfully but the DUT starts self-repairing
-  # (b/130786578), so it's necessary to do sanity check.
-  if not is_good_dut(host):
-    raise errors.ExternalError(
-        'although cros flash succeeded, the DUT is in bad state')
+  return True
 
 
-def cros_flash_with_retry(chromeos_root,
-                          host,
-                          board,
-                          image_path,
-                          version=None,
-                          clobber_stateful=False,
-                          disable_rootfs_verification=True,
-                          repair_callback=None,
-                          force_reboot_callback=None):
+def provision_image_with_retry(chromeos_root,
+                               host,
+                               board,
+                               image_info,
+                               version=None,
+                               clobber_stateful=False,
+                               disable_rootfs_verification=True,
+                               repair_callback=None,
+                               force_reboot_callback=None):
   # 'cros flash' is not 100% reliable, retry if necessary.
   for attempt in range(2):
     if attempt > 0:
@@ -1038,11 +1184,11 @@
       time.sleep(60)
 
     try:
-      cros_flash(
+      provision_image(
           chromeos_root,
           host,
           board,
-          image_path,
+          image_info,
           version=version,
           clobber_stateful=clobber_stateful,
           disable_rootfs_verification=disable_rootfs_verification,
@@ -1307,6 +1453,25 @@
   shutil.copy(src, dst_outside)
 
 
+def _copy_template_files(src, dst):
+  if not os.path.exists(src):
+    return
+
+  def copy_if_nonexistent(src, dst):
+    if not os.path.exists(dst):
+      shutil.copy2(src, dst)
+
+  shutil.copytree(
+      src, dst, dirs_exist_ok=True, copy_function=copy_if_nonexistent)
+
+
+def override_autotest_config(autotest_dir):
+  shadow_config_path = os.path.join(autotest_dir, 'shadow_config.ini')
+  if not os.path.exists(shadow_config_path):
+    with open(shadow_config_path, 'w') as f:
+      f.write(autotest_shadow_config)
+
+
 def prepare_chroot(chromeos_root):
   mount_chroot(chromeos_root)
 
@@ -1317,6 +1482,16 @@
   if os.path.exists(os.path.expanduser(creds_path)):
     copy_into_chroot(chromeos_root, creds_path, creds_path, overwrite=False)
 
+  # quick-provision requires special config for autotest.
+  override_autotest_config(os.path.join(chromeos_root, in_tree_autotest_dir))
+
+  # Copy optional configure files into the home directory inside chromeos
+  # chroot. For example, quick-provision may need special ssh config.
+  assert os.environ.get('USER')
+  _copy_template_files(
+      os.path.join(common.BISECT_KIT_ROOT, 'cros_template_files', 'at_home'),
+      os.path.join(chromeos_root, 'chroot', 'home', os.environ['USER']))
+
 
 def check_if_need_recreate_chroot(stdout, stderr):
   """Analyze build log and determine if chroot should be recreated.
@@ -1488,11 +1663,10 @@
     # For other failures, don't know how to handle. Just bail out.
     raise
 
-  image_symlink = os.path.join(chromeos_root, cached_images_dir, board,
-                               'latest')
+  image_symlink = os.path.join(chromeos_root, build_images_dir, board, 'latest')
   assert os.path.exists(image_symlink)
   image_name = os.readlink(image_symlink)
-  image_folder = os.path.join(cached_images_dir, board, image_name)
+  image_folder = os.path.join(build_images_dir, board, image_name)
   assert os.path.exists(
       os.path.join(chromeos_root, image_folder, test_image_filename))
   return image_folder
diff --git a/eval_cros_autotest.py b/eval_cros_autotest.py
index bbe8403..aa820ff 100755
--- a/eval_cros_autotest.py
+++ b/eval_cros_autotest.py
@@ -17,7 +17,6 @@
 import re
 import subprocess
 import sys
-import textwrap
 
 from bisect_kit import catapult_util
 from bisect_kit import cli
@@ -175,7 +174,8 @@
                                 cros_util.prebuilt_autotest_dir)
   else:
     autotest_dir = os.path.join(opts.chromeos_root,
-                                'src/third_party/autotest/files')
+                                cros_util.in_tree_autotest_dir)
+  cros_util.override_autotest_config(autotest_dir)
   sox_path = os.path.join(opts.chromeos_root, 'chroot/usr/bin/sox')
   if not os.path.exists(sox_path):
     try:
@@ -185,17 +185,6 @@
       # chromeos (b/136136270), so ignore the failure.
       logger.debug('Sox is only required by some audio tests. '
                    'Assume the failure of installing sox is harmless')
-  shadow_config_path = os.path.join(autotest_dir, 'shadow_config.ini')
-  if not os.path.exists(shadow_config_path):
-    with open(shadow_config_path, 'w') as f:
-      f.write(
-          textwrap.dedent("""
-            [CROS]
-            enable_ssh_tunnel_for_servo: True
-            enable_ssh_tunnel_for_chameleon: True
-            enable_ssh_connection_for_devserver: True
-            enable_ssh_tunnel_for_moblab: True
-          """))
 
   # test_that may use this ssh key and ssh complains its permission is too open.
   # chmod every time just before run test_that because the permission may change
diff --git a/switch_cros_localbuild.py b/switch_cros_localbuild.py
index 82d32d3..0ecd89d 100755
--- a/switch_cros_localbuild.py
+++ b/switch_cros_localbuild.py
@@ -210,11 +210,14 @@
   if opts.noimage or opts.nodeploy:
     return 0
 
-  if cros_util.cros_flash_with_retry(
+  image_info = cros_util.ImageInfo()
+  image_info[cros_util.ImageType.DISK_IMAGE] = image_path
+
+  if cros_util.provision_image_with_retry(
       opts.chromeos_root,
       opts.dut,
       opts.board,
-      image_path,
+      image_info,
       clobber_stateful=opts.clobber_stateful,
       disable_rootfs_verification=opts.disable_rootfs_verification,
       repair_callback=cros_lab_util.repair,
@@ -276,6 +279,8 @@
   if cros_util.is_cros_short_version(opts.rev):
     opts.rev = cros_util.version_to_full(opts.board, opts.rev)
 
+  cros_util.prepare_chroot(opts.chromeos_root)
+
   try:
     returncode = switch(opts)
   except Exception:
diff --git a/switch_cros_localbuild_buildbucket.py b/switch_cros_localbuild_buildbucket.py
index 8f6471c..e0f504f 100755
--- a/switch_cros_localbuild_buildbucket.py
+++ b/switch_cros_localbuild_buildbucket.py
@@ -24,7 +24,6 @@
 from bisect_kit import errors
 from bisect_kit import gclient_util
 from bisect_kit import repo_util
-from bisect_kit import util
 
 logger = logging.getLogger(__name__)
 
@@ -298,8 +297,7 @@
   return result.to_string()
 
 
-def prepare_image(image_folder, buildbucket_id):
-  image_path = os.path.join(image_folder, cros_util.test_image_filename)
+def search_build_image(buildbucket_id):
   api = buildbucket_util.BuildbucketApi()
   properties = api.get_build(buildbucket_id).output.properties
 
@@ -307,19 +305,13 @@
     raise errors.ExternalError('artifacts not found in buildbucket_id: %s' %
                                buildbucket_id)
 
-  gs_path = 'gs://%s/%s/image.zip' % (properties['artifacts']['gs_bucket'],
-                                      properties['artifacts']['gs_path'])
-  cros_util.gsutil('cp', gs_path, image_folder)
-  util.check_call(
-      'unzip',
-      '-j',
-      'image.zip',
-      cros_util.test_image_filename,
-      cwd=image_folder)
-  os.remove(os.path.join(image_folder, 'image.zip'))
+  gs_path = 'gs://%s/%s' % (properties['artifacts']['gs_bucket'],
+                            properties['artifacts']['gs_path'])
 
-  assert os.path.exists(image_path)
-  return image_path
+  image_info = cros_util.ImageInfo()
+  image_info[cros_util.ImageType.PARTITION_IMAGE] = gs_path
+  image_info[cros_util.ImageType.ZIP_FILE] = gs_path + '/image.zip'
+  return image_info
 
 
 def schedule_build(opts):
@@ -389,39 +381,27 @@
 
 def switch(opts):
   assert opts.board
-  cached_name = 'bisect-%s' % util.escape_rev(opts.chromeos_rev)
-  if opts.subcommand == 'bisect_chrome':
-    cached_name += '-%s' % util.escape_rev(opts.chrome_rev)
-  image_folder = os.path.join(opts.chromeos_root, cros_util.cached_images_dir,
-                              opts.board, cached_name)
-  image_path = os.path.join(image_folder, cros_util.test_image_filename)
-  if not os.path.exists(image_folder):
-    os.makedirs(image_folder)
 
   # schedule build request and prepare image
-  buildbucket_id = None
-  if not os.path.exists(image_path):
-    if opts.buildbucket_id:
-      buildbucket_id = opts.buildbucket_id
-    else:
-      buildbucket_id = schedule_build(opts)
+  if opts.buildbucket_id:
+    buildbucket_id = opts.buildbucket_id
+  else:
+    buildbucket_id = schedule_build(opts)
 
   # If --nodeploy is given, the script will exit as soon as job scheduled and
   # ignore potential build errors.
   if opts.nodeploy:
     return 0
 
-  if buildbucket_id:
-    wait_build_complete(buildbucket_id)
-    image_path = prepare_image(image_folder, buildbucket_id)
-  rel_path = os.path.relpath(image_path, opts.chromeos_root)
+  wait_build_complete(buildbucket_id)
+  image_info = search_build_image(buildbucket_id)
 
   # deploy and flash
-  if cros_util.cros_flash_with_retry(
+  if cros_util.provision_image_with_retry(
       opts.chromeos_root,
       opts.dut,
       opts.board,
-      rel_path,
+      image_info,
       clobber_stateful=opts.clobber_stateful,
       disable_rootfs_verification=opts.disable_rootfs_verification,
       repair_callback=cros_lab_util.repair,
@@ -458,6 +438,8 @@
   if opts.build_revlist:
     build_revlist(opts)
 
+  cros_util.prepare_chroot(opts.chromeos_root)
+
   try:
     returncode = switch(opts)
   except Exception:
diff --git a/switch_cros_prebuilt.py b/switch_cros_prebuilt.py
index 14130e9..744b5c4 100755
--- a/switch_cros_prebuilt.py
+++ b/switch_cros_prebuilt.py
@@ -15,7 +15,6 @@
 from bisect_kit import common
 from bisect_kit import configure
 from bisect_kit import cros_lab_util
-from bisect_kit import cros_quick_provision
 from bisect_kit import cros_util
 
 logger = logging.getLogger(__name__)
@@ -59,38 +58,22 @@
       default=configure.get('DEFAULT_CHROMEOS_ROOT',
                             os.path.expanduser('~/chromiumos')),
       help='Default chromeos tree to run "cros flash" (default: %(default)s)')
-  parser.add_argument(
-      '--use-quick-provision',
-      '--use_quick_provision',
-      action='store_true',
-      help='Runs quick-provision instead of cros flash.')
-  parser.add_argument(
-      '--build',
-      type=cros_quick_provision.argtype_cros_build,
-      help='ChromeOS build to fetch from gs cache with quick-provision.')
 
   return parser
 
 
 def switch(opts):
-  if opts.use_quick_provision:
-    cros_quick_provision.quick_provision(
-        opts.default_chromeos_root, opts.dut, opts.build)
-    return 0
-
   # TODO(kcwu): clear cache of cros flash
-  if cros_util.is_cros_snapshot_version(opts.version):
-    image_path = cros_util.prepare_snapshot_image(opts.default_chromeos_root,
-                                                  opts.board, opts.version)
-  else:
-    image_path = cros_util.prepare_prebuilt_image(opts.default_chromeos_root,
-                                                  opts.board, opts.version)
+  image_info = cros_util.search_image(opts.board, opts.version)
+  if not image_info:
+    logger.error('no images available for %s %s', opts.board, opts.version)
+    return cli.EXIT_CODE_FATAL
 
-  if cros_util.cros_flash_with_retry(
+  if cros_util.provision_image_with_retry(
       opts.default_chromeos_root,
       opts.dut,
       opts.board,
-      image_path,
+      image_info,
       version=opts.version,
       clobber_stateful=opts.clobber_stateful,
       disable_rootfs_verification=opts.disable_rootfs_verification,
@@ -114,6 +97,8 @@
   if not opts.board:
     opts.board = cros_util.query_dut_board(opts.dut)
 
+  cros_util.prepare_chroot(opts.default_chromeos_root)
+
   try:
     returncode = switch(opts)
   except Exception: