cli/lib/core/provision.py

#
# Copyright (C) 2016 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Provision an attached device."""


import os
import shutil
import tempfile

from core import tool
import error


class Error(error.Error):
    """Base provision error class."""


class MissingBuildError(Error):
    """Raised when the required binaries have not yet been built."""


SYSTEM_IMAGE_FILENAME = 'system.img'


class _ProvisionDir(object):
    """Creates a directory to run `provision-device` from.

    This replicates the platform build directory using symlinks, except
    that system.img will be replaced with a link to a custom file if
    one is available.

    This is meant to be used as a context manager so that it's easy to
    be sure the directory is cleaned up before exiting.
    """

    def __init__(self, product_build_out, system_image_dir=None):
        """Creates a temp dir with the necessary files to provision.

        Args:
            product_build_out: product build output directory.
            system_image_dir: directory containing the custom system.img file
                to use. If None or system.img does not exist, uses the one in
                platform_build_out.
        """
        self.temp_dir = tempfile.mkdtemp(prefix='bdk-flash-')

        try:
            # Look for a custom system.img we should use instead.
            custom_image_file = None
            if system_image_dir:
                image_path = os.path.join(system_image_dir,
                                          SYSTEM_IMAGE_FILENAME)
                if os.path.isfile(image_path):
                    custom_image_file = image_path

            self._fill_temp_dir(product_build_out, custom_image_file)
        except:
            # Wipe the directory if anything goes wrong.
            self.cleanup()
            raise

    def _fill_temp_dir(self, product_build_out, custom_image_file):
        """Fills the temporary directory with the required symlinks."""
        # Not everything in product_build_out is necessary, but it's easier to
        # just link it all and provision-device will pick up whatever it needs.
        for filename in os.listdir(product_build_out):
            link_dst = os.path.join(self.temp_dir, filename)

            # Replace the default system.img if we have a custom one.
            if custom_image_file and filename == SYSTEM_IMAGE_FILENAME:
                link_src = custom_image_file
            else:
                link_src = os.path.join(product_build_out, filename)

            os.symlink(link_src, link_dst)

    def __enter__(self):
        return self.temp_dir

    def __exit__(self, exception_type, value, traceback):
        self.cleanup()

    def cleanup(self):
        shutil.rmtree(self.temp_dir, ignore_errors=True)


def _get_provision_tool(platform, provision_dir):
    """Creates a ProvisionDeviceTool with proper settings.

    Args:
        platform: The Platform to get the provision tool for.
        provision_dir: The dir with images to flash.

    Returns:
        The created ProvisionDeviceTool.

    Raises:
        MissingBuildError: the platform has not been built yet.
    """
    # The provision-device script needs fastboot to be on PATH, so we create
    # a fastboot tool to add it's parent directory to the provision tool PATH.
    fastboot_tool = tool.HostToolWrapper('fastboot', platform)
    provision_tool = tool.ProvisionDeviceTool(
        platform, provision_dir,
        env={'PATH': os.path.dirname(fastboot_tool.path())})

    for t in (fastboot_tool, provision_tool):
        if not t.exists():
            raise MissingBuildError(
                '{} does not exist; use `bdk build platform` first'.format(
                    t.path()))

    return provision_tool


def provision_device(platform, system_image_dir=None, provision_args=None):
    """Provisions the attached device using the `provision-device` script.

    Requires that the platform has been built so that the image files,
    fastboot, and provision-device all exist in platform_build_out.

    Args:
        platform: the Platform to provision.
        system_image_dir: directory containing the custom system.img to
            flash. If None or the system.img file does not exist, uses
            the default system.img from platform_build_out instead.
        provision_args: list of string args to pass to provision-device.

    Raises:
        MissingBuildError: the platform has not been built yet.
        core.util.OSVersionError: the target requests an invalid os version.
        core.tool.Error: (or a subclass) a problem occurs running the tool.
    """
    # Combine the platform ANDROID_PRODUCT_OUT directory with our system.img.
    with _ProvisionDir(platform.product_out_cache,
                       system_image_dir) as provision_dir:
        provision_tool = _get_provision_tool(platform, provision_dir)
        provision_tool.run(provision_args)