Setup kubeconfig with eks and aws-vault

setup-kubectl-eks-aws-vault.py

#!/usr/bin/env python3
import json
import os
import shlex
from configparser import ConfigParser
from copy import deepcopy
from pathlib import Path
from subprocess import check_call, check_output

from ruamel import yaml

HOME = Path(os.environ["HOME"])
AWS_CONFIG = HOME / ".aws" / "config"
KUBECONFIG = HOME / ".kube" / "config"


def main():
    config = ConfigParser()
    config.read_file(AWS_CONFIG.open())
    profiles = "\n".join(config.sections())
    selected_profiles = check_output(
        shlex.split("fzf --multi --header='Select the profiles to configure'"), input=profiles.encode()
    ).decode()
    changed_clusters = []
    for selected_profile in selected_profiles.splitlines():
        profile_name = selected_profile.split()[-1]
        clusters = check_output(
            shlex.split(f"aws-vault exec {profile_name} -- aws eks list-clusters --query 'clusters[*]'")
        ).decode()
        parsed_clusters = json.loads(clusters)
        for parsed_cluster in parsed_clusters:
            check_call(
                shlex.split(
                    f"aws-vault exec {profile_name} -- aws eks update-kubeconfig --name {parsed_cluster} --alias {parsed_cluster}"
                )
            )
            changed_clusters.append((parsed_cluster, profile_name))

    patch_command(changed_clusters)


def patch_command(changed_clusters):
    parsed = yaml.safe_load(KUBECONFIG.open())
    old = deepcopy(parsed)
    for user in parsed["users"]:
        for cluster, profile in changed_clusters:
            if (
                user["name"].endswith(f"/{cluster}")
                and "exec" in user["user"]
                and user["user"]["exec"]["command"] == "aws"
            ):
                user["user"]["exec"]["command"] = "aws-vault"
                args = shlex.split(f"exec {profile} -- aws") + user["user"]["exec"]["args"]
                user["user"]["exec"]["args"] = args
    if old != parsed:
        yaml.safe_dump(parsed, KUBECONFIG.open("w"))


if __name__ == "__main__":
    main()