Merge pull request #610 from mplsgrant/2024-09-install-helm

Offer to install helm into venv

Author: m3dwards

src/warnet/constants.py

@@ -107,3 +107,54 @@
     "helm repo update",
     f"helm upgrade --install ingress-nginx ingress-nginx/ingress-nginx --namespace {INGRESS_NAMESPACE} --create-namespace",
 ]
+
+# Helm binary
+HELM_DOWNLOAD_URL_STUB = "https://get.helm.sh/"
+HELM_BINARY_NAME = "helm"
+HELM_BLESSED_VERSION = "v3.16.1"
+HELM_BLESSED_NAME_AND_CHECKSUMS = [
+    {
+        "name": "helm-v3.16.1-darwin-amd64.tar.gz",
+        "checksum": "1b194824e36da3e3889920960a93868b541c7888c905a06757e88666cfb562c9",
+    },
+    {
+        "name": "helm-v3.16.1-darwin-arm64.tar.gz",
+        "checksum": "405a3b13f0e194180f7b84010dfe86689d7703e80612729882ad71e2a4ef3504",
+    },
+    {
+        "name": "helm-v3.16.1-linux-amd64.tar.gz",
+        "checksum": "e57e826410269d72be3113333dbfaac0d8dfdd1b0cc4e9cb08bdf97722731ca9",
+    },
+    {
+        "name": "helm-v3.16.1-linux-arm.tar.gz",
+        "checksum": "a15a8ddfc373628b13cd2a987206756004091a1f6a91c3b9ee8de6f0b1e2ce90",
+    },
+    {
+        "name": "helm-v3.16.1-linux-arm64.tar.gz",
+        "checksum": "780b5b86f0db5546769b3e9f0204713bbdd2f6696dfdaac122fbe7f2f31541d2",
+    },
+    {
+        "name": "helm-v3.16.1-linux-386.tar.gz",
+        "checksum": "92d7a47a90734b50528ffffc99cd1b2d4b9fc0f4291bac92c87ef03406a5a7b2",
+    },
+    {
+        "name": "helm-v3.16.1-linux-ppc64le.tar.gz",
+        "checksum": "9f0178957c94516eff9a3897778edb93d78fab1f76751bd282883f584ea81c23",
+    },
+    {
+        "name": "helm-v3.16.1-linux-s390x.tar.gz",
+        "checksum": "357f8b441cc535240f1b0ba30a42b44571d4c303dab004c9e013697b97160360",
+    },
+    {
+        "name": "helm-v3.16.1-linux-riscv64.tar.gz",
+        "checksum": "9a2cab45b7d9282e9be7b42f86d8034dcaa2e81ab338642884843676c2f6929f",
+    },
+    {
+        "name": "helm-v3.16.1-windows-amd64.zip",
+        "checksum": "89952ea1bace0a9498053606296ea03cf743c48294969dfc731e7f78d1dc809a",
+    },
+    {
+        "name": "helm-v3.16.1-windows-arm64.zip",
+        "checksum": "fc370a291ed926da5e77acf42006de48e7fd5ff94d20c3f6aa10c04fea66e53c",
+    },
+]

src/warnet/project.py

@@ -1,15 +1,26 @@
+import hashlib
 import os
 import platform
+import shutil
 import subprocess
 import sys
+import tarfile
+import tempfile
 from dataclasses import dataclass
 from enum import Enum, auto
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Optional
 
 import click
 import inquirer
+import requests
 
+from .constants import (
+    HELM_BINARY_NAME,
+    HELM_BLESSED_NAME_AND_CHECKSUMS,
+    HELM_BLESSED_VERSION,
+    HELM_DOWNLOAD_URL_STUB,
+)
 from .graph import inquirer_create_network
 from .network import copy_network_defaults, copy_scenario_defaults
 
@@ -155,7 +166,7 @@ def is_kubectl_installed() -> tuple[bool, str]:
         except FileNotFoundError as err:
             return False, str(err)
 
-    def is_helm_installed() -> tuple[bool, str]:
+    def is_helm_installed_and_offer_if_not() -> tuple[bool, str]:
         try:
             version_result = subprocess.run(["helm", "version"], capture_output=True, text=True)
             location_result = subprocess.run(
@@ -167,8 +178,31 @@ def is_helm_installed() -> tuple[bool, str]:
                 return version_result.returncode == 0, location_result.stdout.strip()
             else:
                 return False, ""
-        except FileNotFoundError as err:
-            return False, str(err)
+
+        except FileNotFoundError:
+            print()
+            helm_answer = inquirer.prompt(
+                [
+                    inquirer.Confirm(
+                        "install_helm",
+                        message=click.style(
+                            "Would you like Warnet to install Helm into your virtual environment?",
+                            fg="blue",
+                            bold=True,
+                        ),
+                        default=True,
+                    ),
+                ]
+            )
+            if helm_answer is None:
+                msg = "Setup cancelled by user."
+                click.secho(msg, fg="yellow")
+                return False, msg
+            if helm_answer["install_helm"]:
+                click.secho("    Installing Helm...", fg="yellow", bold=True)
+                install_helm_rootlessly_to_venv()
+                return is_helm_installed_and_offer_if_not()
+            return False, "Please install Helm."
 
     def check_installation(tool_info: ToolInfo) -> ToolStatus:
         has_good_version, location = tool_info.is_installed_func()
@@ -218,8 +252,8 @@ def check_installation(tool_info: ToolInfo) -> ToolStatus:
     )
     helm_info = ToolInfo(
         tool_name="Helm",
-        is_installed_func=is_helm_installed,
-        install_instruction="Install Helm from Helm's official site.",
+        is_installed_func=is_helm_installed_and_offer_if_not,
+        install_instruction="Install Helm from Helm's official site, or rootlessly install Helm using Warnet's downloader when prompted.",
         install_url="https://helm.sh/docs/intro/install/",
     )
     minikube_info = ToolInfo(
@@ -361,3 +395,137 @@ def init():
     """Initialize a warnet project in the current directory"""
     current_dir = Path.cwd()
     new_internal(directory=current_dir, from_init=True)
+
+
+def get_os_name_for_helm() -> Optional[str]:
+    """Return a short operating system name suitable for downloading a helm binary."""
+    uname_sys = platform.system().lower()
+    if "linux" in uname_sys:
+        return "linux"
+    elif uname_sys == "darwin":
+        return "darwin"
+    elif "win" in uname_sys:
+        return "windows"
+    return None
+
+
+def is_in_virtualenv() -> bool:
+    """Check if the user is in a virtual environment."""
+    return hasattr(sys, "real_prefix") or (
+        hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
+    )
+
+
+def download_file(url, destination):
+    click.secho(f"    Downloading {url}", fg="blue")
+    response = requests.get(url, stream=True)
+    if response.status_code == 200:
+        with open(destination, "wb") as f:
+            for chunk in response.iter_content(1024):
+                f.write(chunk)
+    else:
+        raise Exception(f"Failed to download {url} (status code {response.status_code})")
+
+
+def query_arch_from_uname(arch: str) -> Optional[str]:
+    if arch.startswith("armv5"):
+        return "armv5"
+    elif arch.startswith("armv6"):
+        return "armv6"
+    elif arch.startswith("armv7"):
+        return "arm"
+    elif arch == "aarch64" or arch == "arm64":
+        return "arm64"
+    elif arch == "x86":
+        return "386"
+    elif arch == "x86_64":
+        return "amd64"
+    elif arch == "i686" or arch == "i386":
+        return "386"
+    else:
+        return None
+
+
+def write_blessed_checksum(helm_filename: str, dest_path: str):
+    checksum = next(
+        (b["checksum"] for b in HELM_BLESSED_NAME_AND_CHECKSUMS if b["name"] == helm_filename), None
+    )
+    if checksum:
+        with open(dest_path, "w") as f:
+            f.write(checksum)
+    else:
+        click.secho("Could not find a matching helm binary and checksum", fg="red")
+
+
+def verify_checksum(file_path, checksum_path):
+    click.secho("    Verifying checksum...", fg="blue")
+    sha256_hash = hashlib.sha256()
+    with open(file_path, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+
+    with open(checksum_path) as f:
+        expected_checksum = f.read().strip()
+
+    if sha256_hash.hexdigest() != expected_checksum:
+        raise Exception("Checksum verification failed!")
+    click.secho("    Checksum verified.", fg="blue")
+
+
+def install_helm_to_venv(helm_bin_path):
+    venv_bin_dir = os.path.join(sys.prefix, "bin")
+    helm_dst_path = os.path.join(venv_bin_dir, HELM_BINARY_NAME)
+    shutil.move(helm_bin_path, helm_dst_path)
+    os.chmod(helm_dst_path, 0o755)
+    click.secho(f"    {HELM_BINARY_NAME} installed into {helm_dst_path}", fg="blue")
+
+
+def install_helm_rootlessly_to_venv():
+    if not is_in_virtualenv():
+        click.secho(
+            "Error: You are not in a virtual environment. Please activate a virtual environment and try again.",
+            fg="yellow",
+        )
+        sys.exit(1)
+
+    version = HELM_BLESSED_VERSION
+
+    os_name = get_os_name_for_helm()
+    if os_name is None:
+        click.secho(
+            "Error: Could not determine the operating system of this computer.", fg="yellow"
+        )
+        sys.exit(1)
+
+    uname_arch = os.uname().machine
+    arch = query_arch_from_uname(uname_arch)
+    if not arch:
+        click.secho(f"No Helm binary candidate for arch: {uname_arch}", fg="red")
+        sys.exit(1)
+
+    helm_filename = f"{HELM_BINARY_NAME}-{version}-{os_name}-{arch}.tar.gz"
+    helm_url = f"{HELM_DOWNLOAD_URL_STUB}{helm_filename}"
+
+    try:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            helm_archive_path = os.path.join(temp_dir, helm_filename)
+            checksum_path = os.path.join(temp_dir, f"{helm_filename}.sha256")
+
+            download_file(helm_url, helm_archive_path)
+            write_blessed_checksum(helm_filename, checksum_path)
+            verify_checksum(helm_archive_path, checksum_path)
+
+            # Extract Helm and install it in the virtual environment's bin folder
+            with tarfile.open(helm_archive_path, "r:gz") as tar:
+                tar.extractall(path=temp_dir)
+            helm_bin_path = os.path.join(temp_dir, os_name + "-" + arch, HELM_BINARY_NAME)
+            install_helm_to_venv(helm_bin_path)
+
+            click.secho(
+                f"    {HELM_BINARY_NAME} {version} installed successfully to your virtual environment!\n",
+                fg="blue",
+            )
+
+    except Exception as e:
+        click.secho(f"Error: {e}\nCould not install helm.", fg="yellow")
+        sys.exit(1)