Merge pull request #1 from apache/register-runner-script
Add script to help store self-hosted runner creds in AWS SSM
diff --git a/.flake8 b/.flake8
new file mode 100644
index 0000000..6e7651c
--- /dev/null
+++ b/.flake8
@@ -0,0 +1,20 @@
+# 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.
+[flake8]
+max-line-length = 110
+ignore = E203,E231,E731,W504,I001,W503
+exclude = .svn,CVS,.bzr,.hg,.git,__pycache__,.eggs,*.egg
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..02a11bd
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,91 @@
+# 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.
+---
+default_stages: [commit, push]
+default_language_version:
+ # force all unspecified python hooks to run python3
+ python: python3
+minimum_pre_commit_version: "1.20.0"
+repos:
+ - repo: meta
+ hooks:
+ - id: identity
+ - id: check-hooks-apply
+ - repo: https://github.com/Lucas-C/pre-commit-hooks
+ rev: v1.1.9
+ hooks:
+ - id: forbid-tabs
+ - id: insert-license
+ name: Add license
+ exclude: ^\.github/.*$|^license-templates/
+ args:
+ - --comment-style
+ - "|#|"
+ - --license-filepath
+ - license-templates/LICENSE.txt
+ - --fuzzy-match-generates-todo
+ - id: insert-license
+ name: Add license for all rst files
+ exclude: ^\.github/.*$
+ args:
+ - --comment-style
+ - "||"
+ - --license-filepath
+ - license-templates/LICENSE.rst
+ - --fuzzy-match-generates-todo
+ files: \.rst$
+ - repo: https://github.com/psf/black
+ rev: 20.8b1
+ hooks:
+ - id: black
+ args: [--config=./pyproject.toml]
+ - repo: https://github.com/pre-commit/pre-commit-hooks
+ rev: v3.4.0
+ hooks:
+ - id: check-merge-conflict
+ - id: debug-statements
+ - id: check-builtin-literals
+ - id: detect-private-key
+ - id: end-of-file-fixer
+ - id: mixed-line-ending
+ - id: trailing-whitespace
+ - id: fix-encoding-pragma
+ args:
+ - --remove
+ - repo: https://github.com/asottile/pyupgrade
+ rev: v2.7.4
+ hooks:
+ - id: pyupgrade
+ args: ["--py36-plus"]
+ - repo: https://github.com/pre-commit/pygrep-hooks
+ rev: v1.7.0
+ hooks:
+ - id: rst-backticks
+ - id: python-no-log-warn
+ - repo: https://github.com/timothycrosley/isort
+ rev: 5.7.0
+ hooks:
+ - id: isort
+ name: Run isort to sort imports
+ files: \.py$
+ # To keep consistent with the global isort skip config defined in setup.cfg
+ exclude: ^build/.*$|^.tox/.*$|^venv/.*$
+ - repo: https://gitlab.com/pycqa/flake8
+ rev: 3.8.4
+ hooks:
+ - id: flake8
+ name: Run flake8
diff --git a/README.rst b/README.rst
index a10db90..8582b74 100644
--- a/README.rst
+++ b/README.rst
@@ -1,2 +1,19 @@
+ .. 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.
+
CI Infrastructure for Apache Airflow
====================================
diff --git a/license-templates/LICENSE.rst b/license-templates/LICENSE.rst
new file mode 100644
index 0000000..adf897d
--- /dev/null
+++ b/license-templates/LICENSE.rst
@@ -0,0 +1,16 @@
+.. 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.
diff --git a/license-templates/LICENSE.txt b/license-templates/LICENSE.txt
new file mode 100644
index 0000000..60b675e
--- /dev/null
+++ b/license-templates/LICENSE.txt
@@ -0,0 +1,16 @@
+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.
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..0333bb0
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,28 @@
+# 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.
+[tool.black]
+line-length = 110
+target-version = ['py36', 'py37', 'py38']
+skip-string-normalization = true
+
+[tool.isort]
+line_length = 110
+combine_as_imports = true
+default_section = 'THIRDPARTY'
+# Need to be consistent with the exclude config defined in pre-commit-config.yaml
+skip = ['build','.tox','venv']
+profile = 'black'
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..432240d
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,20 @@
+# 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.
+
+boto3
+click~=7.1
+requests
diff --git a/scripts/store-agent-creds.py b/scripts/store-agent-creds.py
new file mode 100755
index 0000000..7f00dfa
--- /dev/null
+++ b/scripts/store-agent-creds.py
@@ -0,0 +1,206 @@
+#!/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()