engflow_auth.py

"""
Install engflow_auth binary if not on system and authenticate if token is expired.
"""

import json
import os
import platform
import stat
import subprocess
import time
import urllib.request
from datetime import datetime

from retry import retry

NORMALIZED_ARCH = {"x86_64": "x64", "aarch64": "arm64", "arm64": "arm64", "AMD64": "x64"}

NORMALIZED_OS = {"Windows": "windows", "Darwin": "macos", "Linux": "linux"}

GH_URL_PREFIX = "https://github.com/EngFlow/auth/releases/latest/download/"
CLUSTER = "sodalite.cluster.engflow.com"

LOGIN_TIMEOUT = 300
CLI_WAIT_TIMEOUT = 5


@retry(tries=3, delay=1)
def download(url: str, path: str):
    urllib.request.urlretrieve(url, filename=path)


def get_release_tag():
    tag = "engflow_auth"
    sys = platform.system()
    os_sys = NORMALIZED_OS.get(sys, None)
    machine = platform.machine()
    arch = NORMALIZED_ARCH.get(machine, None)
    if not os_sys or not arch:
        raise Exception(f"{sys} {machine} not supported")

    if os_sys == "windows":
        return tag + "_windows_x64.exe"
    return tag + f"_{os_sys}_{arch}"


def install(verbose: bool) -> str:
    binary_directory = os.path.expanduser("~/.local/bin")
    os.makedirs(binary_directory, exist_ok=True)
    binary_filename = "engflow_auth"
    binary_path = os.path.join(binary_directory, binary_filename)
    tag = get_release_tag()
    if "windows" in tag:
        binary_path += ".exe"
    if os.path.exists(binary_path):
        if verbose:
            print(f"{binary_filename} already exists at {binary_path}, skipping download")
    else:
        url = GH_URL_PREFIX + tag
        print(f"Downloading {url}...")
        download(url, binary_path)
        os.chmod(binary_path, stat.S_IEXEC)
        print(f"Downloaded {binary_path}")

    return binary_path


def update_bazelrc(binary_path: str, verbose: bool):
    norm_path = os.path.normpath(binary_path).replace("\\", "/")
    lines = []
    bazelrc_path = f"{os.path.expanduser('~')}/.bazelrc"
    if verbose:
        print(f"Updating {bazelrc_path}")
    if os.path.exists(bazelrc_path):
        with open(bazelrc_path, "r") as bazelrc:
            for line in bazelrc.readlines():
                if "--tls_client" in line or "--credential_helper" in line:
                    continue
                lines.append(line)
    lines.append(f"common --credential_helper={CLUSTER}={norm_path}")

    with open(bazelrc_path, "w+") as bazelrc:
        bazelrc.writelines(lines)


def authenticate(binary_path: str, verbose: bool) -> bool:
    need_login = False
    p = subprocess.run(f"{binary_path} export {CLUSTER}", shell=True, capture_output=True)
    if p.returncode != 0:
        need_login = True
    else:
        expiry_iso = json.loads(p.stdout)["token"]["expiry"][:23]
        if datetime.now() > datetime.fromisoformat(expiry_iso):
            need_login = True
    if not need_login:
        if verbose:
            print("Already authenticated. Skipping authentication.")
        return True

    p = subprocess.Popen(
        f"{binary_path} login -store=file {CLUSTER}",
        shell=True,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )

    login_url = None

    start_time = time.time()
    while not login_url and time.time() < start_time + CLI_WAIT_TIMEOUT:
        line = p.stderr.readline().strip()
        if line.startswith("https://" + CLUSTER):
            login_url = line
            break

    if not login_url:
        print("CLI had unexpected output.")
        p.kill()
        return False

    print(
        f"On any device with a browser, login via the following link to complete EngFlow authentication:\n{login_url}"
    )

    try:
        p.wait(timeout=LOGIN_TIMEOUT)
        print("Sucessfully authenticated with EngFlow!")
    except subprocess.TimeoutExpired:
        print(
            "Timed out waiting for login attempt. Failed to authenticate with EngFlow. Builds will be run locally..."
        )
        p.kill()
        return False

    return True


def setup_auth(verbose: bool = True) -> bool:
    path = install(verbose)
    authenticated = authenticate(path, verbose)
    if not authenticated:
        return False
    update_bazelrc(path, verbose)
    return True


def main():
    return 0 if setup_auth() else 1


if __name__ == "__main__":
    main()