recipes/contrib/android_artifact_downloader.py

# Copyright 2023 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Downloads Android artifacts from go/ab then uploads to CIPD."""

from collections import defaultdict
from PB.recipe_engine.result import RawResult
from PB.go.chromium.org.luci.buildbucket.proto import common
from PB.recipes.fuchsia.contrib.android_artifact_downloader import InputProperties

DEPS = [
    "fuchsia/buildbucket_util",
    "fuchsia/cipd_ensure",
    "fuchsia/cipd_util",
    "recipe_engine/archive",
    "recipe_engine/cipd",
    "recipe_engine/file",
    "recipe_engine/path",
    "recipe_engine/properties",
    "recipe_engine/raw_io",
    "recipe_engine/step",
]

# Service account used to obtain OAuth for accessing Android build service.
# We use LUCI Auth Impersonation go/luci-authorization#service-accounts-impersonation
# to get authentication for accessing Android build API. This allows us to keep
# separate Cloud projects for using different Google Cloud APIs.
SERVICE_ACCOUNT = (
    "aemu-artifact-mover@fuchsia-cloud-api-for-test.iam.gserviceaccount.com"
)

# Auth scope used with luci-auth. Note that we need the additional scope
# "https://www.googleapis.com/auth/androidbuild.internal" in order to access
# Android build service.
SCOPES = [
    "https://www.googleapis.com/auth/cloud-platform",
    "https://www.googleapis.com/auth/firebase",
    "https://www.googleapis.com/auth/gerritcodereview",
    "https://www.googleapis.com/auth/userinfo.email",
    "https://www.googleapis.com/auth/androidbuild.internal",
]

PROPERTIES = InputProperties


def RunSteps(api, props):
    """Checks latest available build from the target branch containing all
    targets, then build and upload cipd packages if they don't already exist
    for that build id.
    """

    # The android_downloader tool doesn't expect all of the artifact metadata,
    # so this formats them into a dict<string><array<string>> for consumption.
    ad_artifacts = defaultdict(list)
    for artifact_spec in props.artifact_specs:
        ad_artifacts[artifact_spec.target] += [artifact_spec.file_path]

    luci_auth = api.cipd_ensure(
        api.resource("luci_auth/cipd.ensure"),
        "infra/tools/luci-auth/${platform}",
    )
    ab_downloader = api.cipd_ensure(
        api.resource("android_downloader/cipd.ensure"),
        "turquoise_internal/infra/android-build-downloader/${platform}",
    )

    # Check the latest Android BuildID (bid) from go/ab that contains all
    # targets.
    bid = find_bid(api, ad_artifacts, props.branch, luci_auth, ab_downloader)

    cipd_tag = str(bid)
    if props.patch_version:
        cipd_tag = f"{props.patch_version}@{cipd_tag}"

    if not props.dry_run:
        # Search CIPD packages to see if we've built for this bid already.
        found_build = check_if_build_exists(api, cipd_tag, props.cipd_pkg_paths)
        if found_build:
            return RawResult(
                status=common.SUCCESS,
                summary_markdown=f"found builds for bid:{cipd_tag}, exiting early.",
            )

    # Download all artifacts from the build.
    download_root = api.path.cleanup_dir.join(bid)
    api.file.ensure_directory("ensure download dir", download_root)

    args = [
        luci_auth,
        "context",
        "-scopes",
        " ".join(SCOPES),
        "-act-as-service-account",
        SERVICE_ACCOUNT,
        "--",
        ab_downloader,
        "-bid",
        bid,
        "-output",
        download_root,
        "-branch",
        props.branch,
    ]

    for target, paths in ad_artifacts.items():
        for path in paths:
            args.extend(["-artifact", f"{target},{path}"])

    api.step(f"download artifacts from branch: {props.branch} bid: {bid}", args)

    # Apply transformations to downloaded artifacts to shape them for upload
    # to CIPD.
    pkg_to_path = {}
    for cipd_path in props.cipd_pkg_paths:
        pkg_root = api.path.cleanup_dir.join(cipd_path.replace("/", "_"))
        api.file.ensure_directory(f"ensure CIPD dir {pkg_root}", pkg_root)
        pkg_to_path[cipd_path] = pkg_root

    apply_transformations(api, props.artifact_specs, bid, download_root, pkg_to_path)

    for pkg, path in pkg_to_path.items():
        api.file.listdir(
            f"list files to be uploaded for {pkg}",
            path,
            test_data=["foo", "bin/bar", "baz.zip"],
        )

    # Upload packages to CIPD.
    if not props.dry_run:
        upload_packages(api, cipd_tag, pkg_to_path)

    return RawResult(
        status=common.SUCCESS,
        summary_markdown=f"successfully uploaded packages to cipd: {props.cipd_pkg_paths}.",
    )


def find_bid(api, artifacts, branch, luci_auth, ab_downloader):
    with api.step.nest("find latest build containing all artifacts"):
        args = [
            luci_auth,
            "context",
            "-scopes",
            " ".join(SCOPES),
            "-act-as-service-account",
            SERVICE_ACCOUNT,
            "--",
            ab_downloader,
            "-find_bid",
            "-branch",
            branch,
        ]

        for target, paths in artifacts.items():
            for path in paths:
                args.extend(["-artifact", f"{target},{path}"])

        step = api.step(
            f"get latest bid from Android build for {branch}",
            args,
            stdout=api.raw_io.output_text(),
            step_test_data=lambda: api.raw_io.test_api.stream_output_text("1234567"),
        )
        bid = step.stdout.strip()

        # TODO(fxb/125660): move error handling to android_downloader.
        error_msg = None
        if not bid:
            error_msg = "android_downloader did not find a build id to download"
        elif not bid.isdigit():
            error_msg = f"android_downloader returned a non-numerical bid: {bid}"
        if error_msg:
            step.presentation.status = api.step.EXCEPTION
            raise api.step.InfraFailure(error_msg)

        return bid


def check_if_build_exists(api, cipd_tag, cipd_pkg_paths):
    with api.step.nest("check if CIPD packages already exist for this bid"):
        for path in cipd_pkg_paths:
            found_pkgs = api.cipd.search(
                package_name=path,
                tag=f"bid:{cipd_tag}",
                test_instances=None,
            )
            if found_pkgs:
                return True
        return False


def apply_transformations(api, artifacts, bid, download_root, pkg_to_path):
    # TODO(fxb/125660): move transformations to android_downloader.
    """Apply transformations to files to make them suitable for the package.
    This currently supports a minimum set of exclusive transformations:
    extract - unpack from a compressed file
    make_executable - chmod +x
    """
    with api.step.nest("apply transformations to downloaded files"):
        for artifact in artifacts:
            # Populate the placeholder {bid} with the actual bid.
            fp = artifact.file_path.replace("{bid}", str(bid))
            file_src = api.path.join(download_root, artifact.target, fp)

            if artifact.extract:
                ext_root = api.path.cleanup_dir.join(f"{artifact.target}_{fp}")

                api.file.ensure_directory(
                    f"ensure extraction directory {ext_root}", ext_root
                )
                api.archive.extract(
                    f"extract {fp}",
                    archive_file=file_src,
                    output=ext_root,
                    include_files=artifact.include_files,
                )

                if artifact.extract_subdir:
                    ext_root = api.path.abs_to_path(
                        api.path.abspath(ext_root.join(artifact.extract_subdir))
                    )

                with api.step.nest(
                    f"move contents of {ext_root} to {artifact.cipd_path} package"
                ):
                    for f in api.file.listdir(
                        f"list files in {ext_root}",
                        ext_root,
                        test_data=["foo", "bin/bar", "baz.zip"],
                    ):
                        api.file.move(
                            f"move {f} to {artifact.cipd_path} package",
                            f,
                            api.path.join(
                                pkg_to_path[artifact.cipd_path], api.path.basename(f)
                            ),
                        )
            else:
                file_dest = api.path.join(pkg_to_path[artifact.cipd_path], fp)
                api.file.ensure_directory(
                    f"ensure destination {file_dest}", api.path.dirname(file_dest)
                )
                if artifact.make_executable:
                    api.file.chmod(f"make {fp} executable", file_src, "0755")
                api.file.move(
                    f"move {fp} to {artifact.cipd_path} package",
                    file_src,
                    file_dest,
                )


def upload_packages(api, cipd_tag, pkg_to_path):
    with api.step.nest("upload packages to CIPD"):
        for pkg, path in pkg_to_path.items():
            api.cipd_util.upload_package(
                pkg_name=pkg,
                pkg_root=path,
                install_mode="copy",
                search_tag={"bid": cipd_tag},
            )


def GenTests(api):
    def properties(**kwargs):
        defaults = {
            "branch": "android-branch",
            "cipd_pkg_paths": [
                "path/to/pkg/",
                "path/to/pkg2/",
            ],
            "artifact_specs": [
                {
                    "target": "build-target",
                    "cipd_path": "path/to/pkg/",
                    "file_path": "folder.zip",
                    "extract": True,
                    "extract_subdir": "secrets/",
                },
                {
                    "target": "build-target",
                    "cipd_path": "path/to/pkg/",
                    "file_path": "bin/binary",
                    "make_executable": True,
                },
            ],
        }
        return api.properties(**{**defaults, **kwargs})

    yield (
        api.buildbucket_util.test("basic")
        + properties(patch_version=1)
        + api.step_data(
            "check if CIPD packages already exist for this bid.cipd search path/to/pkg/ bid:1@1234567",
            api.cipd.example_search("path/to/pkg", []),
        )
        + api.step_data(
            "check if CIPD packages already exist for this bid.cipd search path/to/pkg2/ bid:1@1234567",
            api.cipd.example_search("path/to/pkg", []),
        )
    )

    yield (api.buildbucket_util.test("dry_run") + properties(dry_run=True))

    yield (api.buildbucket_util.test("pkg_already_exists") + properties())

    yield (
        api.buildbucket_util.test("android_downloader_failed_bid", status="FAILURE")
        + properties()
        + api.step_data(
            "find latest build containing all artifacts.get latest bid from Android build for android-branch",
            retcode=1,
            stdout=api.raw_io.output_text(""),
        )
    )

    yield (
        api.test("android_downloader_non_numeric_bid", status="INFRA_FAILURE")
        + properties()
        + api.step_data(
            "find latest build containing all artifacts.get latest bid from Android build for android-branch",
            stdout=api.raw_io.output_text("deadbeef"),
        )
    )

    yield (
        api.test("android_downloader_empty_bid", status="INFRA_FAILURE")
        + properties()
        + api.step_data(
            "find latest build containing all artifacts.get latest bid from Android build for android-branch",
            stdout=api.raw_io.output_text(""),
        )
    )