[Core] az logout: Remove service principal access tokens from token…

… cache (#29441)

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

@@ -302,16 +302,25 @@ def login_in_cloud_shell(self):
         return deepcopy(consolidated)
 
     def logout(self, user_or_sp):
+        # The order of below steps matter! We must
+        #   1. Remove the account from MSAL token cache and SP store
+        #   2. Remove the account from CLI profile
+        # This way, if step 1 fails, CLI still keeps track of the account. Otherwise, if we do the
+        # reverse and step 1 fails, CLI will lose track of the account.
+
+        # Step 1: Remove the account from MSAL token cache and SP store (SP only)
+        # We can't distinguish whether user_or_sp is a user or SP, so try both
+        identity = _create_identity_instance(self.cli_ctx, self._authority)
+        identity.logout_user(user_or_sp)
+        identity.logout_service_principal(user_or_sp)
+
+        # Step 2: Remove the account from CLI profile
         subscriptions = self.load_cached_subscriptions(all_clouds=True)
         result = [x for x in subscriptions
                   if user_or_sp.lower() == x[_USER_ENTITY][_USER_NAME].lower()]
         subscriptions = [x for x in subscriptions if x not in result]
         self._storage[_SUBSCRIPTIONS] = subscriptions
 
-        identity = _create_identity_instance(self.cli_ctx, self._authority)
-        identity.logout_user(user_or_sp)
-        identity.logout_service_principal(user_or_sp)
-
     def logout_all(self):
         self._storage[_SUBSCRIPTIONS] = []
 

src/azure-cli-core/azure/cli/core/auth/identity.py

@@ -11,7 +11,7 @@
 from azure.cli.core._environment import get_config_dir
 from knack.log import get_logger
 from knack.util import CLIError
-from msal import PublicClientApplication
+from msal import PublicClientApplication, ConfidentialClientApplication
 
 # Service principal entry properties
 from .msal_authentication import _CLIENT_ID, _TENANT, _CLIENT_SECRET, _CERTIFICATE, _CLIENT_ASSERTION, \
@@ -203,8 +203,9 @@ def login_with_managed_identity(self, scopes, identity_id=None):  # pylint: disa
     def login_in_cloud_shell(self, scopes):
         raise NotImplementedError
 
-    def logout_user(self, user):
-        accounts = self._msal_app.get_accounts(user)
+    def logout_user(self, username):
+        # If username is an SP client ID, it is ignored
+        accounts = self._msal_app.get_accounts(username)
         for account in accounts:
             self._msal_app.remove_account(account)
 
@@ -218,12 +219,21 @@ def logout_all_users(self):
         for e in file_extensions.values():
             _try_remove(self._token_cache_file + e)
 
-    def logout_service_principal(self, sp):
-        # remove service principal secrets
-        self._service_principal_store.remove_entry(sp)
+    def logout_service_principal(self, client_id):
+        # If client_id is a username, it is ignored
+
+        # Step 1: Remove SP from MSAL token cache
+        # Note that removing SP access tokens shouldn't rely on SP store
+        cca = ConfidentialClientApplication(client_id, **self._msal_app_kwargs)
+        cca.remove_tokens_for_client()
+
+        # Step 2: Remove SP from SP store
+        self._service_principal_store.remove_entry(client_id)
 
     def logout_all_service_principal(self):
         # remove service principal secrets
+        # TODO: As MSAL provides no interface to get all service principals in its token cache, this method can't
+        #   clear all service principals' access tokens from MSAL token cache.
         for e in file_extensions.values():
             _try_remove(self._secret_file + e)
 

src/azure-cli-core/azure/cli/core/auth/tests/test_identity.py

@@ -117,6 +117,37 @@ def test_login_with_service_principal_certificate_cert_err(self):
             identity.login_with_service_principal("00000000-0000-0000-0000-000000000000",
                                                   {"certificate": test_cert_file}, "openid")
 
+    @mock.patch("msal.application.PublicClientApplication.remove_account")
+    @mock.patch("msal.application.PublicClientApplication.get_accounts")
+    def test_logout_user(self, get_accounts_mock, remove_account_mock):
+        accounts = [
+            {
+                'home_account_id': '00000000-0000-0000-0000-000000000000.00000000-0000-0000-0000-000000000000',
+                'environment': 'login.microsoftonline.com',
+                'username': '[email protected]',
+                'account_source': 'broker',
+                'authority_type': 'MSSTS',
+                'local_account_id': '00000000-0000-0000-0000-000000000000',
+                'realm': '00000000-0000-0000-0000-000000000000'
+            }
+        ]
+        get_accounts_mock.return_value = accounts
+
+        identity = Identity('https://login.microsoftonline.com')
+        identity.logout_user('00000000-0000-0000-0000-000000000000')
+        remove_account_mock.assert_called_with(accounts[0])
+
+    @mock.patch("azure.cli.core.auth.identity.ServicePrincipalStore.remove_entry")
+    @mock.patch("msal.application.ConfidentialClientApplication.remove_tokens_for_client")
+    @mock.patch("msal.application.ConfidentialClientApplication.__init__", return_value=None)
+    def test_logout_service_principal(self, init_mock, remove_tokens_for_client_mock, remove_entry_mock):
+        identity = Identity('https://login.microsoftonline.com')
+        client_id = '00000000-0000-0000-0000-000000000000'
+        identity.logout_service_principal(client_id)
+        assert init_mock.call_args.args[0] == client_id
+        remove_tokens_for_client_mock.assert_called_once()
+        remove_entry_mock.assert_called_with(client_id)
+
 
 class TestServicePrincipalAuth(unittest.TestCase):
 

src/azure-cli-core/azure/cli/core/tests/test_profile.py

@@ -1242,8 +1242,9 @@ def mi_auth_factory(*args, **kwargs):
         with self.assertRaisesRegex(CLIError, 'Cloud Shell'):
             cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id)
 
+    @mock.patch('azure.cli.core.auth.identity.Identity.logout_service_principal')
     @mock.patch('azure.cli.core.auth.identity.Identity.logout_user')
-    def test_logout(self, logout_user_mock):
+    def test_logout(self, logout_user_mock, logout_service_principal_mock):
         cli = DummyCli()
 
         storage_mock = {'subscriptions': []}
@@ -1258,10 +1259,13 @@ def test_logout(self, logout_user_mock):
 
         # verify
         self.assertEqual(0, len(storage_mock['subscriptions']))
+        # Make sure logout is attempted on both account types
         logout_user_mock.assert_called_with(self.user1)
+        logout_service_principal_mock.assert_called_with(self.user1)
 
+    @mock.patch('azure.cli.core.auth.identity.Identity.logout_all_service_principal')
     @mock.patch('azure.cli.core.auth.identity.Identity.logout_all_users')
-    def test_logout_all(self, logout_all_users_mock):
+    def test_logout_all(self, logout_all_users_mock, logout_all_service_principal_mock):
         cli = DummyCli()
         # setup
         storage_mock = {'subscriptions': []}
@@ -1280,7 +1284,9 @@ def test_logout_all(self, logout_all_users_mock):
 
         # verify
         self.assertEqual([], storage_mock['subscriptions'])
+        # Make sure logout is attempted on both account types
         logout_all_users_mock.assert_called_once()
+        logout_all_service_principal_mock.assert_called_once()
 
     @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True)
     @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True)