Updated login functionality

Bobbins228 · Bobbins228 · commit 31485f4fff1e · 2023-06-29T11:55:12.000+01:00
diff --git a/requirements.txt b/requirements.txt
@@ -5,4 +5,3 @@ kubernetes>=25.3.0,<27
 codeflare-torchx==0.6.0.dev0
 cryptography==40.0.2
 executing==1.2.0
-jinja2==3.1.2
diff --git a/src/codeflare_sdk/cluster/auth.py b/src/codeflare_sdk/cluster/auth.py
@@ -20,23 +20,12 @@
 """
 
 import abc
-import pathlib
-from kubernetes import config
-from jinja2 import Environment, FileSystemLoader
 import os
+from kubernetes import client, config
 
 global path_set
 path_set = False
 
-"""
-auth = KubeConfigFileAuthentication(
-            kube_config_path="config"
-        )
-auth.load_kube_config()
-
-
-"""
-
 
 class Authentication(metaclass=abc.ABCMeta):
     """
@@ -83,7 +72,7 @@ def logout(self):
 
 class TokenAuthentication(Authentication):
     """
-    `TokenAuthentication` is a subclass of `Authentication`. It can be used to authenticate to an OpenShift
+    `TokenAuthentication` is a subclass of `Authentication`. It can be used to authenticate to a Kubernetes
     cluster when the user has an API token and the API server address.
     """
 
@@ -93,186 +82,88 @@ def __init__(
         server: str = None,
         skip_tls: bool = False,
         ca_cert_path: str = "/etc/pki/tls/certs/ca-bundle.crt",
-        username: str = "user",
     ):
         """
         Initialize a TokenAuthentication object that requires a value for `token`, the API Token
-        and `server`, the API server address for authenticating to an OpenShift cluster.
+        and `server`, the API server address for authenticating to a Kubernetes cluster.
         """
 
         self.token = token
         self.server = server
         self.skip_tls = skip_tls
         self.ca_cert_path = ca_cert_path
-        self.username = username
 
     def login(self) -> str:
         """
-        This function is used to login to a Kubernetes cluster using the user's API token and API server address.
+        This function is used to log in to a Kubernetes cluster using the user's API token and API server address.
         Depending on the cluster, a user can choose to login in with `--insecure-skip-tls-verify` by setting `skip_tls`
         to `True` or `--certificate-authority` by setting `skip_tls` to false and providing a path to a ca bundle with `ca_cert_path`.
-
-        If a user does not have a Kubernetes config file one is created from a template with the appropriate user functionality
-        and if they do it is updated with new credentials.
         """
-        dir = pathlib.Path(__file__).parent.parent.resolve()
-        home = os.path.expanduser("~")
+        global path_set
+        global api_client
         try:
-            security = "insecure-skip-tls-verify: false"
+            configuration = client.Configuration()
+            configuration.api_key_prefix["authorization"] = "Bearer"
+            configuration.host = self.server
+            configuration.api_key["authorization"] = self.token
             if self.skip_tls == False:
-                security = "certificate-authority: %s" % self.ca_cert_path
+                configuration.ssl_ca_cert = self.ca_cert_path
             else:
-                security = "insecure-skip-tls-verify: true"
-
-            env = Environment(
-                loader=FileSystemLoader(f"{dir}/templates"),
-                trim_blocks=True,
-                lstrip_blocks=True,
-            )
-            template = env.get_template("config.yaml")
-            server = self.server
-            cluster_name = server[8:].replace(".", "-")
-            # If there is no .kube folder it is created.
-            if not os.path.isdir("%s/.kube" % home):
-                os.mkdir("%s/.kube" % home)
-
-            # If a config file exists then it will be updated with new fields and values.
-            if os.path.isfile("%s/.kube/config" % home):
-                file = open(r"%s/.kube/config" % home, "r").readlines()
-                write_file = open(r"%s/.kube/config" % home, "w")
-                existing = False
-                # Check for existing config
-                for line in file:
-                    if self.server in line:
-                        existing = True
-
-                if existing == False:
-                    for line in file:
-                        # All of these fields are given new lines underneath with credentials info.
-                        if "clusters:" in line:
-                            write_file.write(line)
-                            write_file.write(
-                                "- cluster:\n    %(security)s\n    server: %(server)s\n  name: %(cluster)s\n"
-                                % {
-                                    "security": security,
-                                    "server": self.server,
-                                    "cluster": cluster_name,
-                                }
-                            )
-                            continue
-                        if "contexts:" in line:
-                            write_file.write(line)
-                            write_file.write(
-                                "- context:\n    cluster: %(cluster)s\n    namespace: default\n    user: %(user)s/%(cluster)s\n  name: default/%(cluster)s/%(user)s\n"
-                                % {"cluster": cluster_name, "user": self.username}
-                            )
-                            continue
-                        if "current-context:" in line:
-                            write_file.write(
-                                "current-context: default/{}/{}\n".format(
-                                    cluster_name, self.username
-                                )
-                            )
-                            continue
-                        if "users:" in line:
-                            write_file.write(line)
-                            write_file.write(
-                                "- name: {}/{}\n  user:\n    token: {}\n".format(
-                                    self.username, cluster_name, self.token
-                                )
-                            )
-                            continue
-
-                        write_file.write(line)
-                else:
-                    # If there is an existing config just update the token and username
-                    for line in file:
-                        if "users:" in line:
-                            write_file.write(line)
-                            write_file.write(
-                                "- name: {}/{}\n  user:\n    token: {}\n".format(
-                                    self.username, cluster_name, self.token
-                                )
-                            )
-                            continue
-                        write_file.write(line)
+                configuration.verify_ssl = False
+            api_client = client.ApiClient(configuration)
+            path_set = False
+            return "Logged into %s" % self.server
+        except client.ApiException as exception:
+            return exception
 
-                response = "Updated config file at %s/.kube/config" % home
-            else:
-                # Create a new config file from the config template and store it in HOME/.kube
-                file = open("%s/.kube/config" % home, "w")
-                file.write(
-                    template.render(
-                        security=security,
-                        server=server,
-                        cluster=cluster_name,
-                        context_name="default/{}/{}".format(
-                            cluster_name, self.username
-                        ),
-                        current_context="default/{}/{}".format(
-                            cluster_name, self.username
-                        ),
-                        username="{}/{}".format(self.username, cluster_name),
-                        token=self.token,
-                    )
-                )
-                response = (
-                    "Logged in and created new config file at %s/.kube/config" % home
-                )
-        except:
-            response = "Error logging in. Have you inputted correct credentials?"
-        return response
+    def api_config_handler():
+        """
+        This function is used to load the api client if the user has logged in
+        """
+        if api_client != None and path_set == False:
+            return api_client
+        else:
+            return None
 
     def logout(self) -> str:
         """
         This function is used to logout of a Kubernetes cluster.
         """
-        home = os.path.expanduser("~")
-        file = open(r"%s/.kube/config" % home, "r")
-        lines = file.readlines()
-        line_count = 0
-        for line in lines:
-            if (
-                "- name: {}/{}".format(self.username, self.server[8:].replace(".", "-"))
-                not in line.strip()
-            ):
-                line_count = line_count + 1
-            else:
-                break
-        # The name, user and token are removed from the config file
-        with open(r"%s/.kube/config" % home, "w") as file:
-            for number, line in enumerate(lines):
-                if number not in [line_count, line_count + 1, line_count + 2]:
-                    file.write(line)
-        print("logged out of user %s" % self.username)
+        global path_set
+        path_set = False
+        global api_client
+        api_client = None
 
 
 class KubeConfigFileAuthentication(KubeConfiguration):
     """
     A class that defines the necessary methods for passing a user's own Kubernetes config file.
-    Specifically this class defines the `load_kube_config()`, `config_check()` and `remove_config()` functions.
+    Specifically this class defines the `load_kube_config()` and `config_check()` functions.
     """
 
     def __init__(self, kube_config_path: str = None):
         self.kube_config_path = kube_config_path
 
     def load_kube_config(self):
+        """
+        Function for loading a user's own predefined Kubernetes config file.
+        """
         global path_set
+        global api_client
         try:
             path_set = True
-            print("Loaded user config file at path %s" % self.kube_config_path)
-            response = config.load_kube_config(self.kube_config_path)
+            api_client = None
+            config.load_kube_config(self.kube_config_path)
+            response = "Loaded user config file at path %s" % self.kube_config_path
         except config.ConfigException:
             path_set = False
             raise Exception("Please specify a config file path")
         return response
 
     def config_check():
-        if path_set == False:
+        """
+        Function for loading the config file at the default config location ~/.kube/config if the user has not
+        specified their own config file or has logged in with their token and server.
+        """
+        if path_set == False and api_client == None:
             config.load_kube_config()
-
-    def remove_config(self) -> str:
-        global path_set
-        path_set = False
-        os.remove(self.kube_config_path)
-        print("Removed config file")
diff --git a/src/codeflare_sdk/cluster/awload.py b/src/codeflare_sdk/cluster/awload.py
@@ -25,7 +25,7 @@
 
 from kubernetes import client, config
 from .cluster import _kube_api_error_handling
-from .auth import KubeConfigFileAuthentication
+from .auth import KubeConfigFileAuthentication, TokenAuthentication
 
 
 class AWManager:
@@ -60,7 +60,9 @@ def submit(self) -> None:
         """
         try:
             KubeConfigFileAuthentication.config_check()
-            api_instance = client.CustomObjectsApi()
+            api_instance = client.CustomObjectsApi(
+                TokenAuthentication.api_config_handler()
+            )
             api_instance.create_namespaced_custom_object(
                 group="mcad.ibm.com",
                 version="v1beta1",
@@ -85,7 +87,9 @@ def remove(self) -> None:
 
         try:
             KubeConfigFileAuthentication.config_check()
-            api_instance = client.CustomObjectsApi()
+            api_instance = client.CustomObjectsApi(
+                TokenAuthentication.api_config_handler()
+            )
             api_instance.delete_namespaced_custom_object(
                 group="mcad.ibm.com",
                 version="v1beta1",
diff --git a/src/codeflare_sdk/cluster/cluster.py b/src/codeflare_sdk/cluster/cluster.py
@@ -24,7 +24,7 @@
 
 from ray.job_submission import JobSubmissionClient
 
-from .auth import KubeConfigFileAuthentication
+from .auth import KubeConfigFileAuthentication, TokenAuthentication
 from ..utils import pretty_print
 from ..utils.generate_yaml import generate_appwrapper
 from .config import ClusterConfiguration
@@ -37,7 +37,6 @@
 )
 
 from kubernetes import client, config
-
 import yaml
 import executing
 
@@ -116,7 +115,9 @@ def up(self):
         namespace = self.config.namespace
         try:
             KubeConfigFileAuthentication.config_check()
-            api_instance = client.CustomObjectsApi()
+            api_instance = client.CustomObjectsApi(
+                TokenAuthentication.api_config_handler()
+            )
             with open(self.app_wrapper_yaml) as f:
                 aw = yaml.load(f, Loader=yaml.FullLoader)
             api_instance.create_namespaced_custom_object(
@@ -137,7 +138,9 @@ def down(self):
         namespace = self.config.namespace
         try:
             KubeConfigFileAuthentication.config_check()
-            api_instance = client.CustomObjectsApi()
+            api_instance = client.CustomObjectsApi(
+                TokenAuthentication.api_config_handler()
+            )
             api_instance.delete_namespaced_custom_object(
                 group="mcad.ibm.com",
                 version="v1beta1",
@@ -249,7 +252,9 @@ def cluster_dashboard_uri(self) -> str:
         """
         try:
             KubeConfigFileAuthentication.config_check()
-            api_instance = client.CustomObjectsApi()
+            api_instance = client.CustomObjectsApi(
+                TokenAuthentication.api_config_handler()
+            )
             routes = api_instance.list_namespaced_custom_object(
                 group="route.openshift.io",
                 version="v1",
@@ -366,7 +371,7 @@ def _kube_api_error_handling(e: Exception):  # pragma: no cover
 def _app_wrapper_status(name, namespace="default") -> Optional[AppWrapper]:
     try:
         KubeConfigFileAuthentication.config_check()
-        api_instance = client.CustomObjectsApi()
+        api_instance = client.CustomObjectsApi(TokenAuthentication.api_config_handler())
         aws = api_instance.list_namespaced_custom_object(
             group="mcad.ibm.com",
             version="v1beta1",
@@ -385,7 +390,7 @@ def _app_wrapper_status(name, namespace="default") -> Optional[AppWrapper]:
 def _ray_cluster_status(name, namespace="default") -> Optional[RayCluster]:
     try:
         KubeConfigFileAuthentication.config_check()
-        api_instance = client.CustomObjectsApi()
+        api_instance = client.CustomObjectsApi(TokenAuthentication.api_config_handler())
         rcs = api_instance.list_namespaced_custom_object(
             group="ray.io",
             version="v1alpha1",
@@ -405,7 +410,7 @@ def _get_ray_clusters(namespace="default") -> List[RayCluster]:
     list_of_clusters = []
     try:
         KubeConfigFileAuthentication.config_check()
-        api_instance = client.CustomObjectsApi()
+        api_instance = client.CustomObjectsApi(TokenAuthentication.api_config_handler())
         rcs = api_instance.list_namespaced_custom_object(
             group="ray.io",
             version="v1alpha1",
@@ -427,7 +432,7 @@ def _get_app_wrappers(
 
     try:
         KubeConfigFileAuthentication.config_check()
-        api_instance = client.CustomObjectsApi()
+        api_instance = client.CustomObjectsApi(TokenAuthentication.api_config_handler())
         aws = api_instance.list_namespaced_custom_object(
             group="mcad.ibm.com",
             version="v1beta1",
@@ -454,7 +459,7 @@ def _map_to_ray_cluster(rc) -> Optional[RayCluster]:
         status = RayClusterStatus.UNKNOWN
 
     KubeConfigFileAuthentication.config_check()
-    api_instance = client.CustomObjectsApi()
+    api_instance = client.CustomObjectsApi(TokenAuthentication.api_config_handler())
     routes = api_instance.list_namespaced_custom_object(
         group="route.openshift.io",
         version="v1",
diff --git a/src/codeflare_sdk/templates/kubconfig-template.yaml b/src/codeflare_sdk/templates/kubconfig-template.yaml
diff --git a/src/codeflare_sdk/utils/generate_cert.py b/src/codeflare_sdk/utils/generate_cert.py
