feat(eks): Allow helm pull from OCI repositories

packages/@aws-cdk/aws-eks/lib/kubectl-handler/helm/__init__.py

@@ -1,8 +1,10 @@
 import json
 import logging
 import os
+import re
 import subprocess
 import shutil
+import tempfile
 import zipfile
 from urllib.parse import urlparse, unquote
 
@@ -78,13 +80,72 @@ def helm_handler(event, context):
             # future work: support versions from s3 assets
             chart = get_chart_asset_from_url(chart_asset_url)
 
+        if repository.startswith('oci://'):
+            assert(repository is not None)
+            tmpdir = tempfile.TemporaryDirectory()
+            chart_dir = get_chart_from_oci(tmpdir.name, release, repository, version)
+            chart = chart_dir
+
         helm('upgrade', release, chart, repository, values_file, namespace, version, wait, timeout, create_namespace)
     elif request_type == "Delete":
         try:
             helm('uninstall', release, namespace=namespace, timeout=timeout)
         except Exception as e:
             logger.info("delete error: %s" % e)
 
+
+def get_oci_cmd(repository, version):
+
+    cmnd = []
+    pattern = '\d+.dkr.ecr.[a-z]+-[a-z]+-\d.amazonaws.com'
+
+    registry = repository.rsplit('/', 1)[0].replace('oci://', '')
+
+    if re.fullmatch(pattern, registry) is not None:
+        region = registry.replace('.amazonaws.com', '').split('.')[-1]
+        cmnd = [
+            "HELM_EXPERIMENTAL_OCI=1;" \
+            f"aws ecr get-login-password --region {region} | " \
+            f"helm registry login --username AWS --password-stdin {registry}; helm pull {repository} --version {version} --untar"
+            ]
+    else:
+        logger.info("Non AWS OCI repository found")
+        cmnd = ['HELM_EXPERIMENTAL_OCI=1', 'helm', 'pull', repository, '--version', version, '--untar']
+
+    return cmnd
+
+
+def get_chart_from_oci(tmpdir, release, repository = None, version = None):
+
+    cmnd = get_oci_cmd(repository, version)
+
+    maxAttempts = 3
+    retry = maxAttempts
+    while retry > 0:
+        try:
+            logger.info(cmnd)
+            env = get_env_with_oci_flag()
+            output = subprocess.check_output(cmnd, stderr=subprocess.STDOUT, cwd=tmpdir, env=env, shell=True)
+            logger.info(output)
+
+            return os.path.join(tmpdir, release)
+        except subprocess.CalledProcessError as exc:
+            output = exc.output
+            if b'Broken pipe' in output:
+                retry = retry - 1
+                logger.info("Broken pipe, retries left: %s" % retry)
+            else:
+                raise Exception(output)
+    raise Exception(f'Operation failed after {maxAttempts} attempts: {output}')
+
+
+def get_env_with_oci_flag():
+    env = os.environ.copy()
+    env['HELM_EXPERIMENTAL_OCI'] = '1'
+
+    return env
+
+
 def helm(verb, release, chart = None, repo = None, file = None, namespace = None, version = None, wait = False, timeout = None, create_namespace = None):
     import subprocess
 
@@ -113,7 +174,8 @@ def helm(verb, release, chart = None, repo = None, file = None, namespace = None
     retry = maxAttempts
     while retry > 0:
         try:
-            output = subprocess.check_output(cmnd, stderr=subprocess.STDOUT, cwd=outdir)
+            env = get_env_with_oci_flag()
+            output = subprocess.check_output(cmnd, stderr=subprocess.STDOUT, cwd=outdir, env=env)
             logger.info(output)
             return
         except subprocess.CalledProcessError as exc: