#!/usr/bin/env python3
# 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/ (if necessary) and then executes it,
passing through all arguments.
"""

import base64
import hashlib
import http.client
import os
import platform
import ssl
import subprocess
import sys
import urllib.parse

# Generated from the following command. May need to be periodically rerun.
# $ cipd ls infra/tools/cipd | perl -pe 's[.*/][];s/^/    "/;s/\s*$/",\n/;'
SUPPORTED_PLATFORMS = (
    "aix-ppc64",
    "linux-386",
    "linux-amd64",
    "linux-arm64",
    "linux-armv6l",
    "linux-mips64",
    "linux-mips64le",
    "linux-mipsle",
    "linux-ppc64",
    "linux-ppc64le",
    "linux-riscv64",
    "linux-s390x",
    "mac-amd64",
    "mac-arm64",
    "windows-386",
    "windows-amd64",
    "windows-arm64",
)


class UnsupportedPlatform(Exception):
    pass


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", "aarch")):
        return machine.replace("aarch", "arm")
    if machine.endswith("64"):
        return "amd64"
    if machine.endswith("86"):
        return "386"
    raise Exception("unrecognized arch: {}".format(machine))


def user_agent():
    """Generate a user-agent based on the project name and current hash."""

    try:
        rev = subprocess.check_output(
            ["git", "-C", SCRIPT_DIR, "rev-parse", "HEAD"]
        ).strip()
    except subprocess.CalledProcessError:
        rev = "???"

    if isinstance(rev, bytes):
        rev = rev.decode()

    return "fuchsia-infra/tools/{}".format(rev)


def actual_hash(path):
    """Hash the file at path and return it."""

    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."""

    expected_plat = "{}-{}".format(platform_normalized(), arch_normalized())

    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 == expected_plat:
                return hashval
    raise Exception("platform {} not in {}".format(expected_plat, DIGESTS_FILE))


def https_connect_with_proxy(target_url):
    """Create HTTPSConnection with proxy support."""

    proxy_env = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
    if proxy_env in (None, ""):
        conn = http.client.HTTPSConnection(target_url)
        return conn

    url = urllib.parse.urlparse(proxy_env)
    conn = http.client.HTTPSConnection(url.hostname, url.port)
    headers = {}
    if url.username and url.password:
        auth = "%s:%s" % (url.username, url.password)
        py_version = sys.version_info.major
        if py_version >= 3:
            headers["Proxy-Authorization"] = "Basic " + str(
                base64.b64encode(auth.encode()).decode()
            )
        else:
            headers["Proxy-Authorization"] = "Basic " + base64.b64encode(auth)
    conn.set_tunnel(target_url, 443, headers)
    return conn


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()

    try:
        conn = https_connect_with_proxy(CIPD_HOST)
    except AttributeError:
        print("=" * 70)
        print(
            """
It looks like this version of Python does not support SSL. This is common
when using Homebrew. If using Homebrew please run the following commands.
If not using Homebrew check how your version of Python was built.

brew install openssl  # Probably already installed, but good to confirm.
brew uninstall python && brew install python
""".strip()
        )
        print("=" * 70)
        raise

    full_platform = "{}-{}".format(platform_normalized(), arch_normalized())
    if full_platform not in SUPPORTED_PLATFORMS:
        raise UnsupportedPlatform(full_platform)

    path = "/client?platform={}&version={}".format(full_platform, version)

    for _ in range(10):
        try:
            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()
        except ssl.SSLError:
            print(
                "\n"
                "SSL error in Python when downloading CIPD client.\n"
                "If using system Python try\n"
                "\n"
                "    sudo pip install certifi\n"
                "\n"
                "And if on the system Python on a Mac try\n"
                "\n"
                "    /Applications/Python 3.6/Install Certificates.command\n"
                "\n"
                "If using Homebrew Python try\n"
                "\n"
                "    brew install openssl\n"
                "    brew uninstall python\n"
                "    brew install python\n"
                "\n"
                "If those don't work, address all the potential issues shown \n"
                "by the following command.\n"
                "\n"
                "    brew doctor\n"
                "\n"
                "Otherwise, check that your machine's Python can use SSL, "
                "testing with the http.client module.",
                file=sys.stderr,
            )
            raise

        # Found client bytes.
        if res.status == http.client.OK:  # pylint: disable=no-else-return
            return content

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

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

    raise Exception(
        "failed to download client from https://{}{}".format(CIPD_HOST, path)
    )


def bootstrap(client, silent=False):
    """Bootstrap cipd client installation."""

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

    if not silent:
        print(
            "Bootstrapping cipd client for {}-{}".format(
                platform_normalized(), arch_normalized()
            )
        )

    tmp_path = client + ".tmp"
    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, 0o755)
    os.rename(tmp_path, client)


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

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


def init(client=CLIENT, silent=False):
    """Install/update cipd client."""

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

    if not os.path.isfile(client):
        bootstrap(client, silent)

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

    return client


def main(client=CLIENT, silent=False):
    """Install/update cipd client."""

    try:
        init(client=client, silent=silent)

    except UnsupportedPlatform:
        # Don't show help message below for this exception.
        raise

    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

    return client


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