Merge pull request #34 from jupyter-naas/33-feat-automatic-credentials-if-jupyterhub_api_token

33 feat automatic credentials if jupyterhub api token

Author: Dr0p42

naas_python/utils/domains_base/authorization.py

@@ -187,6 +187,7 @@ def __init__(
         self,
         port=38745,
         trade_url="https://auth.naas.ai/bearer/workspace/longlived",
+        trade_jupyterhub_url="https://auth.naas.ai/bearer/jupyterhub/longlived",
         timeout=60,  # 60,
         login_url=os.environ.get('NAAS_LOGIN_URL', "https://naas.ai?cli_token=generate_token"),
     ):
@@ -197,6 +198,7 @@ def __init__(
         # TODO: Fix bug where the async call during dev mode, only finishes after the exact timeout period
         self._timeout = timeout
         self.trade_url = trade_url
+        self.trade_jupyterhub_url=trade_jupyterhub_url
         self._login_url = login_url
         self._redirect_uri = None
 
@@ -336,6 +338,7 @@ def access_token(self):
     def _gather_env_credentials(self):
         _SCHEMA_VARS = {
             "NAAS_CREDENTIALS_JWT_TOKEN": "jwt_token",
+            "JUPYTERHUB_API_TOKEN" : "jupyterhub_api_token"
         }
 
         credentials = {}
@@ -356,72 +359,63 @@ def _gather_file_credentials(self, credentials_file_path: Path):
         self._file_contents = json.loads(self._file_contents)
         return self._file_contents
 
-    def _generate_credential_file(self, credentials_file_path: Path):
+    def _generate_credential_file(self, credentials_file_path: Path, jupyterhub_access_token=None, prompt=True):
         # Trade access token (and authenticate) and store the long-lived token in the credentials file
+        if jupyterhub_access_token is not None:
+            access_token = self.trade_for_long_lived_token(access_token=jupyterhub_access_token, access_token_type="jupyterhub")
+            self._jwt_token = access_token
+        else:
+            access_token = self.trade_for_long_lived_token(self.access_token())
+            self._jwt_token = access_token
+            
         try:
-            self._jwt_token = self.trade_for_long_lived_token(self.access_token())
-
             # Create target directory in case it does not exists.
-            os.makedirs(
-                '/'.join(
-                    credentials_file_path.as_posix().split('/')[:-1]
-                ),
-                exist_ok=True
-            )
+            os.makedirs(os.path.join(os.path.dirname(credentials_file_path)), exist_ok=True)
 
             with open(credentials_file_path, "w") as file:
-                file.write(json.dumps({"jwt_token": self._jwt_token}))
+                file.write(json.dumps({"jwt_token": access_token}))
 
-            print(f'\n\t✅ CLI Token successfuly generated and stored to {credentials_file_path.as_posix()}\n\n')
+            if prompt:
+                print(f'\n\t✅ CLI Token successfuly generated and stored to {credentials_file_path.as_posix()}\n\n')
 
-            return self._jwt_token
+            return access_token
         except TimeoutException as e:
             print(f'\n\n\t❌ The process was not able to complete in time. Please try again.\n\n')
             return None
 
-
-
     def check_credentials(self):
         # This method is responsible for inspecting the naas credentials files
         # locally and retrieving the jwt token if it exists and is valid.
-        # The order of priority is as follows:
-        # 1. Check if file exists, if not call appropriate method to start authentication process
-        # 2.a Check if file is empty, if so call appropriate method to start authentication process
-        # 2.b If file is not empty, check if token is valid, if not call appropriate method to start authentication process
-        # 3. If file exists and is not empty and token is valid, set the jwt token to the class property and return
 
         credentials_path = Path(os.path.expanduser("~/.naas/credentials"))
-
-        if credentials_path.exists() and credentials_path.is_file():
-            # First order option for credential gathering is to check the file contents and grab the token
+        _var_credentials = self._gather_env_credentials()
+        
+        # look for the existence of overriding environment variable to create the new file
+        if "jwt_token" in _var_credentials:
+            # Validate stored token is valid... then assign value and return
+            self._jwt_token = _var_credentials["jwt_token"]
+            return self._jwt_token
+        
+        elif "jupyterhub_api_token" in _var_credentials:
+            access_token = _var_credentials["jupyterhub_api_token"]
+            self._generate_credential_file(credentials_file_path=credentials_path, jupyterhub_access_token=access_token, prompt=False)
+            
+        elif credentials_path.exists() and credentials_path.is_file():
+            # Check the file contents and grab the token
             credentials = self._gather_file_credentials(credentials_path)
             logging.debug(f"Credentials file found and not empty.")
-            credentials.update(self._gather_env_credentials())
-
-            # If environment variable is present, override file contents
-            if "jwt_token" in credentials:
-                # Validate stored token is valid... then assign value and return
-                self._jwt_token = credentials["jwt_token"]
-                return self._jwt_token
-
-        else:
-            # As a second order option, look for the existence of overriding environment variable to create the new file
-            _var_credentials = self._gather_env_credentials()
-
-            if "jwt_token" in _var_credentials:
-                # Validate stored token is valid... then assign value and return
-                self._jwt_token = _var_credentials["jwt_token"]
-                return self._jwt_token
-
+            # credentials.update(self._gather_env_credentials())
+            self._jwt_token = credentials["jwt_token"]
+            
+        else :   
             # We could not find any credentials, so we need to start the authentication process
             self._generate_credential_file(credentials_file_path=credentials_path)
 
-        # if not self._jwt_token:
-        #     raise Exception("Could not find any credentials to authenticate with.")
-
-    def trade_for_long_lived_token(self, access_token):
-        # headers = {"Authorization": f"Bearer {access_token}"}
-        url = f"{self.trade_url}/?token={access_token}"
+    def trade_for_long_lived_token(self, access_token, access_token_type="workspace"):
+        if access_token_type == "workspace":
+            url = f"{self.trade_url}/?token={access_token}"
+        if access_token_type == "jupyterhub":
+            url = f"{self.trade_jupyterhub_url}/?token={access_token}"
 
         response = requests.get(url)