Using msal token cache to fetch refresh token inside proxy command for az cli version >= 2.30

src/connectedk8s/HISTORY.rst

@@ -5,6 +5,7 @@ Release History
 1.2.0
 ++++++
 
+* Using MSAL from az-cli version >= 2.30.0
 * Updated CSP version to 1.3.017131
 * Updated GA SDK to 2021-10-01
 * Updated CSP endpoint to CDN

src/connectedk8s/azext_connectedk8s/_constants.py

@@ -73,6 +73,7 @@
 Remove_Config_Fault_Type = "Error while removing old csp config"
 Load_Creds_Fault_Type = "Error while loading accessToken.json"
 Creds_NotFound_Fault_Type = "Credentials of user not found"
+MSAL_cache_not_retrieved_Fault_Type = "MSAL cache could not be retrieved"
 Create_Config_Fault_Type = "Error while creating config file for proxy"
 Run_RefreshThread_Fault_Type = "Error while starting refresh thread"
 Load_Kubeconfig_Fault_Type = "Error while loading kubeconfig"
@@ -108,6 +109,7 @@
 Error_enabling_Features = 'Error while updating agents for enabling features. Please run \"kubectl get pods -n azure-arc\" to check the pods in case of timeout error. Error: {}'
 Error_disabling_Features = 'Error while updating agents for disabling features. Please run \"kubectl get pods -n azure-arc\" to check the pods in case of timeout error. Error: {}'
 Proxy_Kubeconfig_During_Deletion_Fault_Type = 'Encountered proxy kubeconfig during deletion.'
+AZ_CLI_ADAL_TO_MSAL_MIGRATE_VERSION = '2.30.0'
 Cannot_Create_ClusterRoleBindings_Fault_Type = 'Cannot create cluster role bindings on this Kubernets cluster'
 CC_Provider_Namespace_Not_Registered_Fault_Type = "Connected Cluster Provider MS.K8 namespace not registered"
 Default_Namespace_Does_Not_Exist_Fault_Type = "The default namespace defined in the kubeconfig doesn't exist on the kubernetes cluster."

src/connectedk8s/azext_connectedk8s/_utils.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from logging import exception
 import os
 import shutil
 import subprocess
@@ -27,6 +28,8 @@
 import azext_connectedk8s._constants as consts
 from kubernetes import client as kube_client
 from azure.cli.core.azclierror import CLIInternalError, ClientRequestError, ArgumentUsageError, ManualInterrupt, AzureResponseError, AzureInternalError, ValidationError
+from azure.cli.core import get_default_cli
+from packaging import version
 
 logger = get_logger(__name__)
 
@@ -398,6 +401,29 @@ def names(self, names):
         logger.debug("Error while trying to monkey patch the fix for list_node(): {}".format(str(ex)))
 
 
+def use_msal_cache():
+        response_cli_version = az_cli("version --output json")
+        try:
+            cli_version = response_cli_version['azure-cli']
+        except Exception as ex:
+            raise CLIInternalError("Unable to decode the az cli version installed: {}".format(str(ex)))
+        if version.parse(cli_version) >= version.parse(consts.AZ_CLI_ADAL_TO_MSAL_MIGRATE_VERSION):
+            return True
+        else:
+            return False
+
+
+def az_cli(args_str):
+    args = args_str.split()
+    cli = get_default_cli()
+    cli.invoke(args, out_file=open(os.devnull, 'w'))
+    if cli.result.result:
+        return cli.result.result
+    elif cli.result.error:
+        raise cli.result.error
+    return True
+
+
 def check_provider_registrations(cli_ctx):
     try:
         rp_client = _resource_providers_client(cli_ctx)

src/connectedk8s/azext_connectedk8s/custom.py

@@ -14,6 +14,7 @@
 import stat
 import platform
 from azure.core.exceptions import ClientAuthenticationError
+from msal_extensions.token_cache import PersistedTokenCache
 import yaml
 import requests
 import urllib.request
@@ -25,6 +26,7 @@
 from knack.prompting import prompt_y_n
 from knack.prompting import NoTTYException
 from azure.cli.core.commands.client_factory import get_subscription_id
+from msal import PublicClientApplication, ConfidentialClientApplication
 from azure.cli.core._profile import Profile
 from azure.cli.core.util import sdk_no_wait
 from azure.cli.core import telemetry
@@ -40,9 +42,11 @@
 from azext_connectedk8s._client_factory import get_graph_client_service_principals
 import azext_connectedk8s._constants as consts
 import azext_connectedk8s._utils as utils
+from azext_connectedk8s._utils import az_cli
 from glob import glob
 from .vendored_sdks.models import ConnectedCluster, ConnectedClusterIdentity, ListClusterUserCredentialProperties
 from threading import Timer, Thread
+import msal_extensions
 import sys
 import hashlib
 import re
@@ -1635,12 +1639,16 @@ def client_side_proxy_wrapper(cmd,
         requestUri = f'{consts.CSP_Storage_Url}/{consts.RELEASE_DATE_WINDOWS}/arcProxy{operating_system}{consts.CLIENT_PROXY_VERSION}.exe'
         older_version_string = f'.clientproxy\\arcProxy{operating_system}*.exe'
         creds_string = r'.azure\accessTokens.json'
+        msal_token_cache_user = r'.azure\msal_token_cache.bin'
+        msal_token_cache_spn = r'.azure\service_principal_entries.bin'
 
     elif(operating_system == 'Linux' or operating_system == 'Darwin'):
         install_location_string = f'.clientproxy/arcProxy{operating_system}{consts.CLIENT_PROXY_VERSION}'
         requestUri = f'{consts.CSP_Storage_Url}/{consts.RELEASE_DATE_LINUX}/arcProxy{operating_system}{consts.CLIENT_PROXY_VERSION}'
         older_version_string = f'.clientproxy/arcProxy{operating_system}*'
         creds_string = r'.azure/accessTokens.json'
+        msal_token_cache_user = r'.azure/msal_token_cache.json'
+        msal_token_cache_spn = r'.azure/service_principal_entries.json'
 
     else:
         telemetry.set_exception(exception='Unsupported OS', fault_type=consts.Unsupported_Fault_Type,
@@ -1738,26 +1746,68 @@ def client_side_proxy_wrapper(cmd,
             raise FileOperationError("Failed to load credentials." + str(e))
 
         user_name = account['user']['name']
-
-        if user_type == 'user':
-            key = 'userId'
-            key2 = 'refreshToken'
+        use_msal_cache = utils.use_msal_cache()
+        if not use_msal_cache:
+            if user_type == 'user':
+                key = 'userId'
+                key2 = 'refreshToken'
+            else:
+                key = 'servicePrincipalId'
+                key2 = 'accessToken'
+
+            for i in range(len(creds_list)):
+                creds_obj = creds_list[i]
+
+                if key in creds_obj and creds_obj[key] == user_name:
+                    creds = creds_obj[key2]
+                    break
+            if creds == '':
+                telemetry.set_exception(exception='Credentials of user not found.', fault_type=consts.Creds_NotFound_Fault_Type,
+                                        summary='Unable to find creds of user while using ADAL')
+                raise UnclassifiedUserFault("Credentials of user not found.")
         else:
-            key = 'servicePrincipalId'
-            key2 = 'accessToken'
-
-        for i in range(len(creds_list)):
-            creds_obj = creds_list[i]
-
-            if key in creds_obj and creds_obj[key] == user_name:
-                creds = creds_obj[key2]
-                break
-
-        if creds == '':
-            telemetry.set_exception(exception='Credentials of user not found.', fault_type=consts.Creds_NotFound_Fault_Type,
-                                    summary='Unable to find creds of user')
-            raise UnclassifiedUserFault("Credentials of user not found.")
-
+            if user_type == "user":
+                response_user_objectid = az_cli("ad signed-in-user show --query objectId -o tsv")
+                token_cache_location = os.path.expanduser(os.path.join('~', msal_token_cache_user))
+                try:
+                   if operating_system == 'Windows':
+                        persistence = msal_extensions.FilePersistenceWithDataProtection(token_cache_location)
+                   elif operating_system == 'Darwin':
+                        persistence = msal_extensions.KeychainPersistence(token_cache_location)
+                   else:
+                        persistence = msal_extensions.FilePersistence(token_cache_location)
+                   token_cache = msal_extensions.PersistedTokenCache(persistence)      
+                except  Exception as e:
+                    telemetry.set_exception(exception=e, fault_type=consts.MSAL_cache_not_retrieved_Fault_Type,
+                                    summary='Unable to retrieve MSAL cache')
+                    raise CLIInternalError("Failed to authenticate with the cluster")
+                token_cache._reload_if_necessary()
+                home_account_id = response_user_objectid + "." + tenantId
+                owned_by_home_account = {"home_account_id": home_account_id}
+                creds_list = token_cache.find(PersistedTokenCache.CredentialType.REFRESH_TOKEN, query=owned_by_home_account)
+                creds = creds_list[0]['secret']
+            else:
+                token_cache_location = os.path.expanduser(os.path.join('~', msal_token_cache_spn))
+                try:
+                   if operating_system == 'Windows':
+                        persistence = msal_extensions.FilePersistenceWithDataProtection(token_cache_location)
+                   elif operating_system == 'Darwin':
+                        persistence = msal_extensions.KeychainPersistence(token_cache_location)
+                   else:
+                        persistence = msal_extensions.FilePersistence(token_cache_location)
+                   token_cache = msal_extensions.PersistedTokenCache(persistence)      
+                except  Exception as e:
+                    telemetry.set_exception(exception=e, fault_type=consts.MSAL_cache_not_retrieved_Fault_Type,
+                                    summary='Unable to retrieve MSAL cache')
+                    raise CLIInternalError("Failed to authenticate with the cluster")
+                token_cache._reload_if_necessary()
+                token_cache_decrypted = token_cache.serialize()
+                creds_list = json.loads(token_cache_decrypted)
+                creds = creds_list[0]['client_secret']
+            if creds == '':
+                telemetry.set_exception(exception='Credentials of user not found.', fault_type=consts.Creds_NotFound_Fault_Type,
+                                        summary='Unable to find creds of user while using MSAL')
+                raise UnclassifiedUserFault("Credentials of user not found.")
         if user_type != 'user':
             dict_file['identity']['clientSecret'] = creds
     else:

src/connectedk8s/setup.py

@@ -17,7 +17,7 @@
 # TODO: Confirm this is the right version number you want and it matches your
 # HISTORY.rst entry.
 
-VERSION = '1.2.0'
+VERSION = '1.2.1'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers