You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2021/01/12 11:59:05 UTC
[airflow-ci-infra] 01/01: Add script to help store self-hosted
runner creds in AWS SSM
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch register-runner-script
in repository https://gitbox.apache.org/repos/asf/airflow-ci-infra.git
commit 5d12e05c23387991be582a5239aa142902075aab
Author: Ash Berlin-Taylor <as...@firemirror.com>
AuthorDate: Tue Jan 12 11:54:51 2021 +0000
Add script to help store self-hosted runner creds in AWS SSM
We can't create self-hosted runners "on-demand", so we need to
pre-create a "pool" of them for use by the auto-scaled nodes.
This script automated the process of converting the short-lived token in
to long-lived credentials (by using the runner binaries in a temporary
directory) and then storing the resulting files in AWS's ParameterStore
---
.flake8 | 4 +
.pre-commit-config.yaml | 91 ++++++++++++++++++++
license-templates/LICENSE.rst | 16 ++++
license-templates/LICENSE.txt | 16 ++++
pyproject.toml | 12 +++
requirements.txt | 20 +++++
scripts/store-agent-creds.py | 193 ++++++++++++++++++++++++++++++++++++++++++
7 files changed, 352 insertions(+)
diff --git a/.flake8 b/.flake8
new file mode 100644
index 0000000..eba5688
--- /dev/null
+++ b/.flake8
@@ -0,0 +1,4 @@
+[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/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..6153b08
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,12 @@
+[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 100644
index 0000000..c512314
--- /dev/null
+++ b/scripts/store-agent-creds.py
@@ -0,0 +1,193 @@
+# 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 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("--token", help="GitHub runner registration token", required=False)
+@click.option("--index", type=int, required=False)
+def main(token, runner_version, repo, index: int):
+ check_aws_config()
+ dir = make_runner_dir(runner_version)
+
+ if not token:
+ token = click.prompt("GitHub runner registration token")
+
+ if not index:
+ index = get_next_index()
+ click.echo(f"Registering as runner {index}")
+
+ register_runner(dir.name, token, repo, 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() -> int:
+ """Find the next available index to store the runner credentials in AWS SSM ParameterStore"""
+ paginator = boto3.client("ssm").get_paginator("describe_parameters")
+
+ pages = paginator.paginate(
+ ParameterFilters=[{"Key": "Path", "Option": "Recursive", "Values": ["/runners/"]}]
+ )
+
+ 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))
+
+ print(seen)
+ # 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, index: int):
+ os.chdir(dir)
+
+ res = subprocess.call(
+ [
+ "./config.sh",
+ "--unattended",
+ "--url",
+ f"https://github.com/{repo}",
+ "--token",
+ token,
+ "--name",
+ f"Airflow Runner {index}",
+ ]
+ )
+
+ if res != 0:
+ exit(res)
+ _put_runner_creds(index)
+
+
+def _put_runner_creds(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/{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/{index}/credentials", Type="String", Value=fh.read())
+
+ with open(".credentials_rsaparams", encoding='utf-8-sig') as fh:
+ client.put_parameter(Name=f"/runners/{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()