scripts/store-agent-creds.py

#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
import json
import os
import platform
import subprocess
import tempfile
from typing import Optional, Tuple

import boto3
import click
import requests
from botocore.exceptions import NoCredentialsError


@click.command()
@click.option(
    "--runner-version",
    default="2.275.1",
    help="Runner version to register with",
    metavar="VER",
)
@click.option("--repo", default="apache/airflow")
@click.option("--store-as", default="apache/airflow")
@click.option("--runnergroup")
@click.option("--token", help="GitHub runner registration token", required=False)
@click.option("--index", type=int, required=False)
def main(
    token, runner_version, store_as: Optional[str], repo, runnergroup: Optional[str], index: Optional[int]
):
    check_aws_config()
    dir = make_runner_dir(runner_version)

    if not token:
        token = click.prompt("GitHub runner registration token")

    if store_as is None:
        store_as = repo

    if index is None:
        index = get_next_index(store_as)
    click.echo(f"Registering as runner {index}")

    register_runner(dir.name, token, repo, runnergroup, store_as, index)


def check_aws_config():
    click.echo("Checking AWS account credentials")
    try:
        whoami = boto3.client("sts").get_caller_identity()
    except NoCredentialsError:
        click.echo("No AWS credentials found -- maybe you need to set AWS_PROFILE?", err=True)
        exit(1)

    if whoami["Account"] != "827901512104":
        click.echo("Wrong AWS account in use -- maybe you need to set AWS_PROFILE?", err=True)
        exit(1)


def make_runner_dir(version):
    """Extract the runner tar to a temporary directory"""
    dir = tempfile.TemporaryDirectory()

    tar = _get_runner_tar(version)

    subprocess.check_call(
        ["tar", "-xzf", tar],
        cwd=dir.name,
    )

    return dir


def get_next_index(repo: str) -> int:
    """Find the next available index to store the runner credentials in AWS SSM ParameterStore"""
    paginator = boto3.client("ssm").get_paginator("describe_parameters")

    path = os.path.join('/runners/', repo, '')

    pages = paginator.paginate(ParameterFilters=[{"Key": "Path", "Option": "Recursive", "Values": [path]}])

    seen = set()

    for page in pages:
        for param in page['Parameters']:
            name = param['Name']

            # '/runners/1/config' -> '1'
            index = os.path.basename(os.path.dirname(name))
            seen.add(int(index))

    if not seen:
        return 1

    # Fill in any gaps too.
    for n in range(1, max(seen) + 2):
        if n not in seen:
            return n


def register_runner(dir: str, token: str, repo: str, runnergroup: Optional[str], store_as: str, index: int):
    os.chdir(dir)

    cmd = [
        "./config.sh",
        "--unattended",
        "--url",
        f"https://github.com/{repo}",
        "--token",
        token,
        "--name",
        f"Airflow Runner {index}",
    ]

    if runnergroup:
        cmd += ['--runnergroup', runnergroup]

    res = subprocess.call(cmd)

    if res != 0:
        exit(res)
    _put_runner_creds(store_as, index)


def _put_runner_creds(repo: str, index: int):
    client = boto3.client("ssm")

    with open(".runner", encoding='utf-8-sig') as fh:
        # We want to adjust the config before storing it!
        config = json.load(fh)
        config["pullRequestSecurity"] = {}

        client.put_parameter(
            Name=f"/runners/{repo}/{index}/config",
            Type="String",
            Value=json.dumps(config, indent=2),
        )

    with open(".credentials", encoding='utf-8-sig') as fh:
        client.put_parameter(Name=f"/runners/{repo}/{index}/credentials", Type="String", Value=fh.read())

    with open(".credentials_rsaparams", encoding='utf-8-sig') as fh:
        client.put_parameter(Name=f"/runners/{repo}/{index}/rsaparams", Type="SecureString", Value=fh.read())


def _get_system_arch() -> Tuple[str, str]:
    uname = platform.uname()
    if uname.system == "Linux":
        system = "linux"
    elif uname.system == "Darwin":
        system = "osx"
    else:
        raise RuntimeError("Un-supported platform")

    if uname.machine == "x86_64":
        arch = "x64"
    else:
        raise RuntimeError("Un-supported architecture")

    return system, arch


def _get_runner_tar(version) -> str:
    system, arch = _get_system_arch()

    cache = os.path.abspath(".cache")

    try:
        os.mkdir(cache)
    except FileExistsError:
        pass

    fname = f"actions-runner-{system}-{arch}-{version}.tar.gz"
    local_file = os.path.join(cache, fname)

    if os.path.exists(local_file):
        return local_file

    url = f"https://github.com/actions/runner/releases/download/v{version}/{fname}"
    click.echo(f"Getting {url}")
    resp = requests.get(url, stream=True)
    resp.raise_for_status()
    with open(local_file, "wb") as fh, click.progressbar(length=int(resp.headers["content-length"])) as bar:
        for chunk in resp.iter_content(chunk_size=40960):
            fh.write(chunk)
            bar.update(len(chunk))
    return local_file


if __name__ == "__main__":
    main()