[Identity] CAE b3

src/azure-cli-core/azure/cli/core/_identity.py

@@ -17,7 +17,8 @@
     ClientSecretCredential,
     CertificateCredential,
     ManagedIdentityCredential,
-    EnvironmentCredential
+    EnvironmentCredential,
+    TokenCachePersistenceOptions
 )
 
 from ._environment import get_config_dir
@@ -65,6 +66,7 @@ def __init__(self, authority=None, tenant_id=None, client_id=None, **kwargs):
         self._msal_app_instance = None
         # Store for Service principal credential persistence
         self._msal_secret_store = MsalSecretStore(fallback_to_plaintext=self.allow_unencrypted)
+        self._cache_persistence_options = TokenCachePersistenceOptions(name="azcli", allow_unencrypted_storage=True)
 
         # TODO: Allow disabling SSL verification
         # The underlying requests lib of MSAL has been patched with Azure Core by MsalTransportAdapter
@@ -86,13 +88,19 @@ def __init__(self, authority=None, tenant_id=None, client_id=None, **kwargs):
         #   - Access token
         #   - Service principal secret
         #   - Service principal certificate
-        # self._credential_kwargs['logging_enable'] = True
+        self._credential_kwargs['logging_enable'] = True
+
+        # Make MSAL remove existing accounts on successful login.
+        # self._credential_kwargs['remove_existing_account'] = True
+        # from azure.cli.core._msal_patch import patch_token_cache_add
+        # patch_token_cache_add(self.msal_app.remove_account)
 
     def _load_msal_cache(self):
         # sdk/identity/azure-identity/azure/identity/_internal/msal_credentials.py:95
-        from azure.identity._internal.persistent_cache import load_user_cache
+        from azure.identity._persistent_cache import _load_persistent_cache
         # Store for user token persistence
-        cache = load_user_cache(self.allow_unencrypted)
+        cache = _load_persistent_cache(self._cache_persistence_options)
+        cache._reload_if_necessary()  # pylint: disable=protected-access
         return cache
 
     def _build_persistent_msal_app(self, authority):
@@ -104,7 +112,7 @@ def _build_persistent_msal_app(self, authority):
         return msal_app
 
     @property
-    def _msal_app(self):
+    def msal_app(self):
         if not self._msal_app_instance:
             # Build the authority in MSAL style, like https://login.microsoftonline.com/your_tenant
             msal_authority = "{}/{}".format(self.authority, self.tenant_id)
@@ -120,8 +128,7 @@ def login_with_interactive_browser(self, scopes=None):
         credential = InteractiveBrowserCredential(authority=self.authority,
                                                   tenant_id=self.tenant_id,
                                                   client_id=self.client_id,
-                                                  enable_persistent_cache=True,
-                                                  allow_unencrypted_cache=self.allow_unencrypted,
+                                                  cache_persistence_options=self._cache_persistence_options,
                                                   **self._credential_kwargs)
         auth_record = credential.authenticate(scopes=scopes)
         # todo: remove after ADAL token deprecation
@@ -139,9 +146,8 @@ def prompt_callback(verification_uri, user_code, _):
             credential = DeviceCodeCredential(authority=self.authority,
                                               tenant_id=self.tenant_id,
                                               client_id=self.client_id,
-                                              enable_persistent_cache=True,
                                               prompt_callback=prompt_callback,
-                                              allow_unencrypted_cache=self.allow_unencrypted,
+                                              cache_persistence_options=self._cache_persistence_options,
                                               **self._credential_kwargs)
 
             auth_record = credential.authenticate(scopes=scopes)
@@ -163,8 +169,7 @@ def login_with_username_password(self, username, password, scopes=None):
                                                 client_id=self.client_id,
                                                 username=username,
                                                 password=password,
-                                                enable_persistent_cache=True,
-                                                allow_unencrypted_cache=self.allow_unencrypted,
+                                                cache_persistence_options=self._cache_persistence_options,
                                                 **self._credential_kwargs)
         auth_record = credential.authenticate(scopes=scopes)
 
@@ -305,15 +310,15 @@ def login_in_cloud_shell(self, scopes):
         return credential, cloud_shell_identity_info
 
     def logout_user(self, user):
-        accounts = self._msal_app.get_accounts(user)
+        accounts = self.msal_app.get_accounts(user)
         logger.info('Before account removal:')
         logger.info(json.dumps(accounts))
 
         # `accounts` are the same user in all tenants, log out all of them
         for account in accounts:
-            self._msal_app.remove_account(account)
+            self.msal_app.remove_account(account)
 
-        accounts = self._msal_app.get_accounts(user)
+        accounts = self.msal_app.get_accounts(user)
         logger.info('After account removal:')
         logger.info(json.dumps(accounts))
 
@@ -323,25 +328,25 @@ def logout_sp(self, sp):
 
     def logout_all(self):
         # TODO: Support multi-authority logout
-        accounts = self._msal_app.get_accounts()
+        accounts = self.msal_app.get_accounts()
         logger.info('Before account removal:')
         logger.info(json.dumps(accounts))
 
         for account in accounts:
-            self._msal_app.remove_account(account)
+            self.msal_app.remove_account(account)
 
-        accounts = self._msal_app.get_accounts()
+        accounts = self.msal_app.get_accounts()
         logger.info('After account removal:')
         logger.info(json.dumps(accounts))
         # remove service principal secrets
         self._msal_secret_store.remove_all_cached_creds()
 
     def get_user(self, user=None):
-        accounts = self._msal_app.get_accounts(user) if user else self._msal_app.get_accounts()
+        accounts = self.msal_app.get_accounts(user) if user else self.msal_app.get_accounts()
         return accounts
 
     def get_user_credential(self, username):
-        accounts = self._msal_app.get_accounts(username)
+        accounts = self.msal_app.get_accounts(username)
 
         # TODO: Confirm with MSAL team that username can uniquely identify the account
         if not accounts:
@@ -352,8 +357,7 @@ def get_user_credential(self, username):
         auth_record = AuthenticationRecord(self.tenant_id, self.client_id, self.authority,
                                            account['home_account_id'], username)
         return InteractiveBrowserCredential(authentication_record=auth_record, disable_automatic_authentication=True,
-                                            enable_persistent_cache=True,
-                                            allow_unencrypted_cache=self.allow_unencrypted,
+                                            cache_persistence_options=self._cache_persistence_options,
                                             **self._credential_kwargs)
 
     def get_service_principal_credential(self, client_id, use_cert_sn_issuer):
@@ -427,7 +431,6 @@ def serialize_token_cache(self, path=None):
                        "It contains login information of all logged-in users. Make sure you protect it safely.", path)
 
         cache = self._load_msal_cache()
-        cache._reload_if_necessary()  # pylint: disable=protected-access
         with open(path, "w") as fd:
             fd.write(cache.serialize())
 

src/azure-cli-core/azure/cli/core/_msal_patch.py

@@ -0,0 +1,173 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+"""
+A temporary workaround for MSAL limitation
+https://github.com/AzureAD/microsoft-authentication-library-for-python/issues/335
+
+After a successful sign-in, if the sign-in account already exists in the token
+cache, remove it first along with its tokens to prevent MSAL from returning
+cached access tokens from the previous session that may have been revoked.
+
+Otherwise, MSAL will return revoked access tokens, resulting in 401 failure
+which can't be handled by commands that don't support silent reauth.
+"""
+
+# pylint: skip-file
+# flake8: noqa
+
+import json
+import time
+
+from knack.log import get_logger
+
+from msal.oauth2cli.oauth2 import Client
+from msal.token_cache import decode_id_token, canonicalize, decode_part
+
+logger = get_logger(__name__)
+
+
+def patch_token_cache_add(callback):
+
+    def __add(self, event, now=None):
+        # event typically contains: client_id, scope, token_endpoint,
+        # response, params, data, grant_type
+        environment = realm = None
+        if "token_endpoint" in event:
+            _, environment, realm = canonicalize(event["token_endpoint"])
+        if "environment" in event:  # Always available unless in legacy test cases
+            environment = event["environment"]  # Set by application.py
+        response = event.get("response", {})
+        data = event.get("data", {})
+        access_token = response.get("access_token")
+        refresh_token = response.get("refresh_token")
+        id_token = response.get("id_token")
+        id_token_claims = (
+            decode_id_token(id_token, client_id=event["client_id"])
+            if id_token else {})
+        client_info = {}
+        home_account_id = None  # It would remain None in client_credentials flow
+        if "client_info" in response:  # We asked for it, and AAD will provide it
+            client_info = json.loads(decode_part(response["client_info"]))
+            home_account_id = "{uid}.{utid}".format(**client_info)
+        elif id_token_claims:  # This would be an end user on ADFS-direct scenario
+            client_info["uid"] = id_token_claims.get("sub")
+            home_account_id = id_token_claims.get("sub")
+
+        target = ' '.join(event.get("scope") or [])  # Per schema, we don't sort it
+
+        with self._lock:
+            now = int(time.time() if now is None else now)
+
+            if client_info and not event.get("skip_account_creation"):
+                account = {
+                    "home_account_id": home_account_id,
+                    "environment": environment,
+                    "realm": realm,
+                    "local_account_id": id_token_claims.get(
+                        "oid", id_token_claims.get("sub")),
+                    "username": id_token_claims.get("preferred_username")  # AAD
+                                or id_token_claims.get("upn")  # ADFS 2019
+                                or "",  # The schema does not like null
+                    "authority_type":
+                        self.AuthorityType.ADFS if realm == "adfs"
+                        else self.AuthorityType.MSSTS,
+                    # "client_info": response.get("client_info"),  # Optional
+                }
+
+                logger.debug("Remove existing account %r", account)
+                logger.debug("Calling %r", callback)
+                callback(account)
+                self.modify(self.CredentialType.ACCOUNT, account, account)
+
+            if id_token:
+                idt = {
+                    "credential_type": self.CredentialType.ID_TOKEN,
+                    "secret": id_token,
+                    "home_account_id": home_account_id,
+                    "environment": environment,
+                    "realm": realm,
+                    "client_id": event.get("client_id"),
+                    # "authority": "it is optional",
+                }
+                self.modify(self.CredentialType.ID_TOKEN, idt, idt)
+
+            if access_token:
+                expires_in = int(  # AADv1-like endpoint returns a string
+                    response.get("expires_in", 3599))
+                ext_expires_in = int(  # AADv1-like endpoint returns a string
+                    response.get("ext_expires_in", expires_in))
+                at = {
+                    "credential_type": self.CredentialType.ACCESS_TOKEN,
+                    "secret": access_token,
+                    "home_account_id": home_account_id,
+                    "environment": environment,
+                    "client_id": event.get("client_id"),
+                    "target": target,
+                    "realm": realm,
+                    "token_type": response.get("token_type", "Bearer"),
+                    "cached_at": str(now),  # Schema defines it as a string
+                    "expires_on": str(now + expires_in),  # Same here
+                    "extended_expires_on": str(now + ext_expires_in)  # Same here
+                }
+                if data.get("key_id"):  # It happens in SSH-cert or POP scenario
+                    at["key_id"] = data.get("key_id")
+                if "refresh_in" in response:
+                    refresh_in = response["refresh_in"]  # It is an integer
+                    at["refresh_on"] = str(now + refresh_in)  # Schema wants a string
+                self.modify(self.CredentialType.ACCESS_TOKEN, at, at)
+
+            if refresh_token:
+                rt = {
+                    "credential_type": self.CredentialType.REFRESH_TOKEN,
+                    "secret": refresh_token,
+                    "home_account_id": home_account_id,
+                    "environment": environment,
+                    "client_id": event.get("client_id"),
+                    "target": target,  # Optional per schema though
+                    "last_modification_time": str(now),  # Optional. Schema defines it as a string.
+                }
+                if "foci" in response:
+                    rt["family_id"] = response["foci"]
+                self.modify(self.CredentialType.REFRESH_TOKEN, rt, rt)
+
+            app_metadata = {
+                "client_id": event.get("client_id"),
+                "environment": environment,
+            }
+            if "foci" in response:
+                app_metadata["family_id"] = response.get("foci")
+            self.modify(self.CredentialType.APP_METADATA, app_metadata, app_metadata)
+
+    def obtain_token_by_refresh_token(self, token_item, scope=None,
+                                      rt_getter=lambda token_item: token_item["refresh_token"],
+                                      on_removing_rt=None,
+                                      on_updating_rt=None,
+                                      on_obtaining_tokens=None,
+                                      **kwargs):
+        resp = super(Client, self).obtain_token_by_refresh_token(
+            rt_getter(token_item)
+            if not isinstance(token_item, str) else token_item,
+            scope=scope,
+            also_save_rt=on_updating_rt is False,
+            on_obtaining_tokens=on_obtaining_tokens,
+            **kwargs)
+        if resp.get('error') == 'invalid_grant':
+            (on_removing_rt or self.on_removing_rt)(token_item)  # Discard old RT
+        RT = "refresh_token"
+        if on_updating_rt is not False and RT in resp:
+            (on_updating_rt or self.on_updating_rt)(token_item, resp[RT])
+        return resp
+
+    from unittest.mock import patch
+
+    # Temporary patch for https://github.com/AzureAD/microsoft-authentication-library-for-python/issues/335
+    cm_add = patch('msal.token_cache.TokenCache._TokenCache__add', __add)
+
+    # Temporary patch for https://github.com/AzureAD/microsoft-authentication-library-for-python/pull/339
+    cm_obtain_token_by_refresh_token = patch('msal.oauth2cli.oauth2.Client.obtain_token_by_refresh_token',
+                                             obtain_token_by_refresh_token)
+    cm_add.__enter__()
+    cm_obtain_token_by_refresh_token.__enter__()