[126] Kubectl: remove unnecessary shenanigans for MFA (#183)

* remove unnecessary shenanigans for MFA

* fix MFA mess

* create kube config directory if missing

* update tests

* restoring refresh_aws_credentials decorator approach

* no longer required

* fix decorator + add mfa/sso tests

* back to the context-manager approach

* get correct region from path rather than assuming it

* avoid cwd property by now, it breaks tf command

* unnecessary import

Author: Franr

leverage/_utils.py

@@ -1,14 +1,9 @@
 """
     General use utilities.
 """
-import functools
 from subprocess import run
 from subprocess import PIPE
 
-from click.exceptions import Exit
-
-from leverage import logger
-
 
 def clean_exception_traceback(exception):
     """ Delete special local variables from all frames of an exception's traceback
@@ -67,7 +62,6 @@ class CustomEntryPoint:
     Set a custom entrypoint on the container while entering the context.
     Once outside, return it to its original value.
     """
-
     def __init__(self, container, entrypoint):
         self.container = container
         self.old_entrypoint = container.entrypoint
@@ -80,54 +74,31 @@ def __exit__(self, *args, **kwargs):
         self.container.entrypoint = self.old_entrypoint
 
 
-class EmptyEntryPoint(CustomEntryPoint):
+class AwsCredsEntryPoint(CustomEntryPoint):
     """
-    Force an empty entrypoint. This will let you execute any commands freely.
+    Fetching AWS credentials by setting the SSO/MFA entrypoints.
+    This works as a replacement of _prepare_container.
     """
 
     def __init__(self, container):
-        super(EmptyEntryPoint, self).__init__(container, entrypoint="")
-
-
-def refresh_aws_credentials(func):
-    """
-    Use this decorator in the case you want to make sure you will have fresh tokens to interact with AWS
-    during the execution of your wrapped method.
-    """
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        container = args[0]  # this is the "self" of the method you are decorating; a LeverageContainer instance
-
         if container.sso_enabled:
             container._check_sso_token()
-            auth_method = container.TF_SSO_ENTRYPOINT
+            auth_method = f"{container.TF_SSO_ENTRYPOINT} -- "
         elif container.mfa_enabled:
-            auth_method = container.TF_MFA_ENTRYPOINT
-            # TODO: ask why this was necessary
+            auth_method = f"{container.TF_MFA_ENTRYPOINT} -- "
             container.environment.update({
                 "AWS_SHARED_CREDENTIALS_FILE": container.environment["AWS_SHARED_CREDENTIALS_FILE"].replace("tmp", ".aws"),
                 "AWS_CONFIG_FILE": container.environment["AWS_CONFIG_FILE"].replace("tmp", ".aws"),
             })
         else:
-            # no auth method found: skip the refresh
-            return func(*args, **kwargs)
-
-        logger.info("Fetching  AWS credentials...")
-        with CustomEntryPoint(container, f"{auth_method} -- echo"):
-            # this simple echo "Fetching..." will run the SSO/MFA entrypoints underneath
-            # that takes care of the token refresh
-            exit_code = container._start("Fetching done.")
-            if exit_code:
-                raise Exit(exit_code)
-            if container.mfa_enabled:
-                # we need to revert to the original values, otherwise other tools that rely on awscli, like kubectl
-                # won't find the credentials
-                container.environment.update({
-                    "AWS_SHARED_CREDENTIALS_FILE": container.environment["AWS_SHARED_CREDENTIALS_FILE"].replace(".aws", "tmp"),
-                    "AWS_CONFIG_FILE": container.environment["AWS_CONFIG_FILE"].replace(".aws", "tmp"),
-                })
-
-        # we should have a valid token at this point, now execute the original method
-        return func(*args, **kwargs)
-
-    return wrapper
+            auth_method = ""
+
+        super(AwsCredsEntryPoint, self).__init__(container, entrypoint=auth_method)
+
+    def __exit__(self, *args, **kwargs):
+        super(AwsCredsEntryPoint, self).__exit__(*args, **kwargs)
+        if self.container.mfa_enabled:
+            self.container.environment.update({
+                "AWS_SHARED_CREDENTIALS_FILE": self.container.environment["AWS_SHARED_CREDENTIALS_FILE"].replace(".aws", "tmp"),
+                "AWS_CONFIG_FILE": self.container.environment["AWS_CONFIG_FILE"].replace(".aws", "tmp"),
+            })

leverage/container.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 from pathlib import Path
 from datetime import datetime
 
@@ -21,6 +22,15 @@
 from leverage.path import NotARepositoryError
 from leverage.conf import load as load_env
 
+REGION = (
+    r"(.*)"  # project folder
+    # start region 
+    r"(global|(?:[a-z]{2}-(?:gov-)?"
+    r"(?:central|north|south|east|west|northeast|northwest|southeast|southwest|secret|topsecret)-[1-4]))"
+    # end region
+    r"(.*)"  # layer
+)
+
 
 def get_docker_client():
     """ Attempt to get a Docker client from the environment configuration. Halt application otherwise.
@@ -167,6 +177,18 @@ def backend_tfvars(self):
     def guest_aws_credentials_dir(self):
         return f"/root/tmp/{self.project}"
 
+    @property
+    def region(self):
+        """
+        Return the region of the layer.
+        """
+        if matches := re.match(REGION, self.cwd.as_posix()):
+            # the region (group 1) is between the projects folders (group 0) and the layers (group 2)
+            return matches.groups()[1]
+
+        logger.exception(f"No valid region could be found at: {self.cwd.as_posix()}")
+        raise Exit(1)
+
     def ensure_image(self):
         """ Make sure the required Docker image is available in the system. If not, pull it from registry. """
         found_image = self.client.api.images(f"{self.image}:{self.image_tag}")
@@ -269,6 +291,20 @@ def run_func(client, container):
 
         return self._run(container, run_func)
 
+    def _start_with_output(self, command="/bin/bash", *args):
+        """
+        Same than _start but also returns the outputs (by dumping the logs) of the container.
+        """
+        container = self._create_container(True, command, *args)
+
+        def run_func(client, container):
+            dockerpty.start(client=client.api, container=container)
+            exit_code = client.api.inspect_container(container)["State"]["ExitCode"]
+            logs = client.api.logs(container).decode("utf-8")
+            return exit_code, logs
+
+        return self._run(container, run_func)
+
     def start(self, command="/bin/sh", *arguments):
         """ Run command with the given arguments in an interactive container.
 

leverage/containers/kubectl.py

@@ -6,7 +6,7 @@
 from docker.types import Mount
 
 from leverage import logger
-from leverage._utils import chain_commands, EmptyEntryPoint, refresh_aws_credentials
+from leverage._utils import chain_commands, AwsCredsEntryPoint
 from leverage.container import TerraformContainer
 
 
@@ -22,52 +22,52 @@ def __init__(self, client):
 
         self.entrypoint = self.KUBECTL_CLI_BINARY
 
-        host_config_path = str(Path.home() / Path(f".kube/{self.project}"))
+        self.host_kubectl_config_dir = Path.home() / Path(f".kube/{self.project}")
+        if not self.host_kubectl_config_dir.exists():
+            # make sure the folder exists before mounting it
+            self.host_kubectl_config_dir.mkdir(parents=True)
+
         self.container_config["host_config"]["Mounts"].append(
             # the container is expecting a file named "config" here
             Mount(
-                source=host_config_path,
+                source=str(self.host_kubectl_config_dir),
                 target=str(self.KUBECTL_CONFIG_PATH),
                 type="bind",
             )
         )
 
-    @refresh_aws_credentials
     def start_shell(self):
-        with EmptyEntryPoint(self):
-            self._start()
+        with AwsCredsEntryPoint(self):
+            self._start("/bin/bash")
 
-    @refresh_aws_credentials
     def configure(self):
         # make sure we are on the cluster layer
         self.check_for_layer_location()
 
         logger.info("Retrieving k8s cluster information...")
-        with EmptyEntryPoint(self):
-            # generate the command that will configure the new cluster
+        # generate the command that will configure the new cluster
+        with AwsCredsEntryPoint(self):
             add_eks_cluster_cmd = self._get_eks_kube_config()
         # and the command that will set the proper ownership on the config file (otherwise the owner will be "root")
         change_owner_cmd = self._change_kube_file_owner_cmd()
         full_cmd = chain_commands([add_eks_cluster_cmd, change_owner_cmd])
 
         logger.info("Configuring context...")
-        with EmptyEntryPoint(self):
-            # we use _start here because in the case of MFA it will ask for the token
+        with AwsCredsEntryPoint(self):
             exit_code = self._start(full_cmd)
         if exit_code:
             raise Exit(exit_code)
 
         logger.info("Done.")
 
     def _get_eks_kube_config(self) -> str:
-        exit_code, output = self._exec(f"{self.TF_BINARY} output")
+        exit_code, output = self._start_with_output(f"{self.TF_BINARY} output -no-color")
         if exit_code:
             logger.error(output)
             raise Exit(exit_code)
 
-        aws_eks_cmd = next(op for op in output.split("\n") if op.startswith("aws eks update-kubeconfig"))
-        # assuming the cluster container is on the primary region
-        return aws_eks_cmd + f" --region {self.common_conf['region_primary']}"
+        aws_eks_cmd = next(op for op in output.split("\r\n") if op.startswith("aws eks update-kubeconfig"))
+        return aws_eks_cmd + f" --region {self.region}"
 
     def _get_user_group_id(self, user_id) -> int:
         user = pwd.getpwuid(user_id)

tests/test_containers/test_kubectl.py

@@ -33,9 +33,9 @@ def kubectl_container(muted_click_context):
 
 
 def test_get_eks_kube_config(kubectl_container):
-    tf_output = "\naws eks update-kubeconfig --name test-cluster --profile test-profile\n"
-    with patch.object(kubectl_container, "_exec", return_value=(0, tf_output)):
-        kubectl_container.common_conf["region_primary"] = "us-east-1"
+    tf_output = "\r\naws eks update-kubeconfig --name test-cluster --profile test-profile\r\n"
+    with patch.object(kubectl_container, "_start_with_output", return_value=(0, tf_output)):
+        kubectl_container.cwd = Path("/project/account/us-east-1/cluster")
         cmd = kubectl_container._get_eks_kube_config()
 
     assert cmd == AWS_EKS_UPDATE_KUBECONFIG
@@ -45,7 +45,7 @@ def test_get_eks_kube_config_tf_output_error(kubectl_container):
     """
     Test that if the TF OUTPUT fails, we get an error back.
     """
-    with patch.object(kubectl_container, "_exec", return_value=(1, "ERROR!")):
+    with patch.object(kubectl_container, "_start_with_output", return_value=(1, "ERROR!")):
         with pytest.raises(Exit):
             kubectl_container._get_eks_kube_config()
 
@@ -85,7 +85,7 @@ def test_start_shell(kubectl_container):
     container_args = kubectl_container.client.api.create_container.call_args[1]
 
     # we want a shell, so -> /bin/bash with no entrypoint
-    assert container_args["command"] == "/bin/sh"
+    assert container_args["command"] == "/bin/bash"
     assert container_args["entrypoint"] == ""
 
     # make sure we are pointing to the AWS credentials
@@ -106,8 +106,53 @@ def test_start_shell(kubectl_container):
 @patch.object(KubeCtlContainer, "check_for_layer_location", Mock())
 # nor terraform
 @patch.object(KubeCtlContainer, "_get_eks_kube_config", Mock(return_value=AWS_EKS_UPDATE_KUBECONFIG))
-def test_configure(kubectl_container, caplog):
+def test_configure(kubectl_container):
     with patch.object(kubectl_container, "_start", return_value=0) as mock_start:
         kubectl_container.configure()
 
     assert mock_start.call_args[0][0] == f'bash -c "{AWS_EKS_UPDATE_KUBECONFIG} && chown 1234:5678 /root/.kube/config"'
+
+
+#####################
+# test auth methods #
+#####################
+
+def test_start_shell_mfa(kubectl_container):
+    """
+    Make sure the command is executed through the proper MFA script.
+    """
+    # container = KubeCtlContainer(docker_client, env_conf=dict(MFA_ENABLED="true", **FAKE_ENV))
+    # container._run = Mock()
+
+    kubectl_container.enable_mfa()
+    # mock the __exit__ of the context manager to avoid the restoration of the values
+    # otherwise the asserts around /.aws/ wouldn't be possible
+    with patch("leverage._utils.AwsCredsEntryPoint.__exit__"):
+        kubectl_container.start_shell()
+        container_args = kubectl_container.client.api.create_container.call_args[1]
+
+    # we want a shell, so -> /bin/bash with no entrypoint
+    assert container_args["command"] == "/bin/bash"
+    assert container_args["entrypoint"] == "/root/scripts/aws-mfa/aws-mfa-entrypoint.sh -- "
+
+    # make sure we are pointing to the right AWS credentials: /.aws/ folder for MFA
+    assert container_args["environment"]["AWS_CONFIG_FILE"] == "/root/.aws/test/config"
+    assert container_args["environment"]["AWS_SHARED_CREDENTIALS_FILE"] == "/root/.aws/test/credentials"
+
+
+def test_start_shell_sso(kubectl_container):
+    """
+    Make sure the command is executed through the proper SSO script.
+    """
+    kubectl_container.enable_sso()
+    kubectl_container._check_sso_token = Mock(return_value=True)
+    kubectl_container.start_shell()
+    container_args = kubectl_container.client.api.create_container.call_args[1]
+
+    # we want a shell, so -> /bin/bash with no entrypoint
+    assert container_args["command"] == "/bin/bash"
+    assert container_args["entrypoint"] == "/root/scripts/aws-sso/aws-sso-entrypoint.sh -- "
+
+    # make sure we are pointing to the right AWS credentials: /tmp/ folder for SSO
+    assert container_args["environment"]["AWS_CONFIG_FILE"] == "/root/tmp/test/config"
+    assert container_args["environment"]["AWS_SHARED_CREDENTIALS_FILE"] == "/root/tmp/test/credentials"