cipd.py

#!/usr/bin/env python
# Copyright 2019 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.
"""Installs and then runs cipd.

This script installs cipd in ./tools/ and then executes it, passing through
all arguments.
"""

from __future__ import print_function

import hashlib
import httplib
import os
import platform
import subprocess
import sys
import urlparse


SCRIPT_DIR = os.path.dirname(__file__)
VERSION_FILE = os.path.join(SCRIPT_DIR, ".cipd_version")
DIGESTS_FILE = VERSION_FILE + ".digests"
# Put CIPD client in tools so that users can easily get it in their PATH.
CLIENT = os.path.join(SCRIPT_DIR, "tools", "cipd")
CIPD_HOST = "chrome-infra-packages.appspot.com"


def platform_normalized():
    """Normalize platform into format expected in CIPD paths."""

    try:
        os_name = platform.system().lower()
        return {"linux": "linux", "mac": "mac", "darwin": "mac", "windows": "windows",}[
            os_name
        ]
    except KeyError:
        raise Exception("unrecognized os: {}".format(os_name))


def arch_normalized():
    """Normalize arch into format expected in CIPD paths."""

    machine = platform.machine()
    if machine.startswith("arm"):
        return machine
    if machine.endswith("64"):
        return "amd64"
    if machine.endswith("86"):
        return "386"
    raise Exception("unrecognized arch: {}".format(machine))


def user_agent():
    try:
        rev = subprocess.check_output(
            ["git", "-C", SCRIPT_DIR, "rev-parse", "HEAD"]
        ).strip()
    except subprocess.CalledProcessError:
        rev = "???"
    return "fuchsia-infra/tools/{}".format(rev)


def actual_hash(path):
    hasher = hashlib.sha256()
    with open(path, "rb") as ins:
        hasher.update(ins.read())
    return hasher.hexdigest()


def expected_hash():
    """Pulls expected hash from digests file."""

    with open(DIGESTS_FILE, "r") as ins:
        for line in ins:
            line = line.strip()
            if line.startswith("#") or not line:
                continue
            plat, hashtype, hashval = line.split()
            if hashtype == "sha256" and plat == "{}-{}".format(
                platform_normalized(), arch_normalized()
            ):
                return hashval
        raise Exception(
            "platform {} not in {}".format(platform_normalized(), DIGESTS_FILE)
        )


def client_bytes():
    """Pull down the CIPD client and return it as a bytes object.

  Often CIPD_HOST returns a 302 FOUND with a pointer to storage.googleapis.com,
  so this needs to handle redirects, but it shouldn't require the initial
  response to be a redirect either.
  """

    with open(VERSION_FILE, "r") as ins:
        version = ins.read().strip()

    conn = httplib.HTTPSConnection(CIPD_HOST)
    path = "/client?platform={platform}-{arch}&version={version}".format(
        platform=platform_normalized(), arch=arch_normalized(), version=version
    )

    for _ in range(10):
        conn.request("GET", path)
        res = conn.getresponse()
        # Have to read the response before making a new request, so make sure
        # we always read it.
        content = res.read()

        # Found client bytes.
        if res.status == httplib.OK:
            return content

        # Redirecting to another location.
        elif res.status == httplib.FOUND:
            location = res.getheader("location")
            url = urlparse.urlparse(location)
            if url.netloc != conn.host:
                conn = httplib.HTTPSConnection(url.netloc)
            path = "{}?{}".format(url.path, url.query)

        # Some kind of error in this response.
        else:
            break

    raise Exception("failed to download client")


def bootstrap():
    """Bootstrap cipd client installation."""

    client_dir = os.path.dirname(CLIENT)
    if not os.path.isdir(client_dir):
        os.makedirs(client_dir)

    print(
        "Bootstrapping cipd client for {}-{}".format(
            platform_normalized(), arch_normalized()
        )
    )
    tmp_path = os.path.join(SCRIPT_DIR, "tools", ".cipd")
    with open(tmp_path, "wb") as tmp:
        tmp.write(client_bytes())

    expected = expected_hash()
    actual = actual_hash(tmp_path)

    if expected != actual:
        raise Exception(
            "digest of downloaded CIPD client is incorrect, check that "
            "digests file is current"
        )

    os.chmod(tmp_path, 0755)
    os.rename(tmp_path, CLIENT)


def selfupdate():
    """Update cipd client."""

    cmd = [
        CLIENT,
        "selfupdate",
        "-version-file",
        VERSION_FILE,
        "-service-url",
        "https://{}".format(CIPD_HOST),
    ]
    subprocess.check_call(cmd)


def init():
    """Install/update cipd client."""

    os.environ["CIPD_HTTP_USER_AGENT_PREFIX"] = user_agent()

    try:
        if not os.path.isfile(CLIENT):
            bootstrap()

        try:
            selfupdate()
        except subprocess.CalledProcessError:
            print(
                "CIPD selfupdate failed. Bootstrapping then retrying...",
                file=sys.stderr,
            )
            bootstrap()
            selfupdate()

    except Exception:
        print(
            "Failed to initialize CIPD. Run "
            "`CIPD_HTTP_USER_AGENT_PREFIX={user_agent}/manual {client} "
            "selfupdate -version-file '{version_file}'` "
            "to diagnose if this is persistent.".format(
                user_agent=user_agent(), client=CLIENT, version_file=VERSION_FILE,
            ),
            file=sys.stderr,
        )
        raise


if __name__ == "__main__":
    init()
    subprocess.check_call([CLIENT] + sys.argv[1:])