kubectlwrapper.py

#!/usr/bin/env python3
import json
import os
import subprocess
import urllib.request
from pathlib import Path
from subprocess import CompletedProcess
from typing import Any, Dict, Mapping
from urllib.request import Request
import sys
import yaml
import shutil

KUBECONFIG: str = os.environ.get("KUBECONFIG", str(Path("~/.kube/config").expanduser()))


def main() -> None:
    kubeconfig_parsed: Dict[str, Any] = {}

    def update_kubeconfig(kcfile: str) -> None:
        with Path(kcfile).open() as kcfd:
            kubeconfig_parsed.update(yaml.load(kcfd, Loader=yaml.SafeLoader))

    if ":" in KUBECONFIG:
        for kcfile in KUBECONFIG.split(":"):
            update_kubeconfig(kcfile)
    else:
        update_kubeconfig(KUBECONFIG)

    # fmt: off
    current_context_name: str = kubeconfig_parsed["current-context"]
    current_context: Mapping[str, Any] = [c for c in kubeconfig_parsed["contexts"] if c["name"] == current_context_name][0]
    context_user: Mapping[str, Any] = [u for u in kubeconfig_parsed["users"] if u["name"] == current_context["context"]["user"]][0]
    context_cluster: str = [c for c in kubeconfig_parsed["clusters"] if c["name"] == current_context["context"]["cluster"]][0]

    token: str = ""
    if "exec" in context_user["user"]:
        result: CompletedProcess = subprocess.run(
            [context_user["user"]["exec"]["command"]] + context_user["user"]["exec"]["args"],
            capture_output=True,
        )
        token = json.loads(result.stdout.decode())["status"]["token"]

    version: str = json.loads(urllib.request.urlopen(Request(
        f"{context_cluster['cluster']['server']}/version", # type: ignore
        headers={"Authorization": f"Bearer {token}"},
    )).read())["gitVersion"]
    # fmt: on

    asdf_plugins: str = subprocess.run(["asdf", "plugin", "list"], capture_output=True).stdout.decode().split("\n")
    if "kubectl" not in asdf_plugins:
        print("asdf kubectl plugin not installed - installing..")
        subprocess.run(["asdf", "plugin", "add", "kubectl"], capture_output=True)

    kubectl_versions: CompletedProcess = subprocess.run(["asdf", "list", "kubectl"], capture_output=True)
    if (
        kubectl_versions.returncode != 0 and kubectl_versions.stderr is not None and "No versions installed" in kubectl_versions.stderr.decode()
        or kubectl_versions.returncode == 0 and version.lstrip("v") not in kubectl_versions.stdout.decode().strip(" ").split("\n")
    ):
        print(f"kubectl {version} not installed - installing...")
        result = subprocess.run(["asdf", "install", "kubectl", version.lstrip("v")], capture_output=True)
        if result.returncode == 0:
            print(f"kubectl {version} installed successfully.")
        else:
            print("error installing kubectl!")
            raise SystemExit(result.stderr)

    set_version_result: CompletedProcess = subprocess.run(["asdf", "global", "kubectl", version.lstrip("v")], capture_output=True)
    if set_version_result.returncode != 0:
        raise SystemExit(set_version_result.stderr.decode().rstrip("\n"))

    reshim_result: CompletedProcess = subprocess.run(["asdf", "reshim"], capture_output=True)
    if reshim_result.returncode != 0:
        raise SystemExit(reshim_result.stderr)

    os.execve(shutil.which("kubectl"), ["kubectl"] + sys.argv[1:], os.environ)

if __name__ == "__main__":
    main()