From 5b0a74d4e9220bbf875269064d7b14a366ecb56d Mon Sep 17 00:00:00 2001 From: jiasli <4003950+jiasli@users.noreply.github.com> Date: Mon, 8 Nov 2021 12:19:12 +0800 Subject: [PATCH 1/2] Use MSAL HTTP cache --- .../azure/cli/core/auth/identity.py | 119 ++++++++++-------- .../azure/cli/core/auth/persistence.py | 1 + .../cli/core/auth/tests/test_identity.py | 30 ++--- .../azure/cli/core/tests/test_profile.py | 27 ++-- src/azure-cli-core/setup.py | 2 +- .../azure/cli/testsdk/patches.py | 23 ++-- src/azure-cli/requirements.py3.Darwin.txt | 2 +- src/azure-cli/requirements.py3.Linux.txt | 2 +- src/azure-cli/requirements.py3.windows.txt | 2 +- 9 files changed, 115 insertions(+), 93 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/auth/identity.py b/src/azure-cli-core/azure/cli/core/auth/identity.py index a5340bcd612..9d8f281e8b5 100644 --- a/src/azure-cli-core/azure/cli/core/auth/identity.py +++ b/src/azure-cli-core/azure/cli/core/auth/identity.py @@ -16,7 +16,7 @@ from .msal_authentication import _CLIENT_ID, _TENANT, _CLIENT_SECRET, _CERTIFICATE, _CLIENT_ASSERTION, \ _USE_CERT_SN_ISSUER from .msal_authentication import UserCredential, ServicePrincipalCredential -from .persistence import load_persisted_token_cache, file_extensions +from .persistence import load_persisted_token_cache, file_extensions, load_secret_store from .util import check_result AZURE_CLI_CLIENT_ID = '04b07795-8ddb-461a-bbee-02f9e1bf7b46' @@ -31,10 +31,19 @@ class Identity: # pylint: disable=too-many-instance-attributes - service principal - TODO: managed identity """ - # HTTP cache for MSAL's tenant discovery, retry-after error cache, etc. - # It must follow singleton pattern. Otherwise, a new dbm.dumb http_cache can read out-of-sync dat and dir. + + # MSAL token cache. + # It follows singleton pattern so that all MSAL app instances share the same token cache. + _msal_token_cache = None + + # MSAL HTTP cache for MSAL's tenant discovery, retry-after error cache, etc. + # It *must* follow singleton pattern so that all MSAL app instances shares the same HTTP cache. # https://github.com/AzureAD/microsoft-authentication-library-for-python/pull/407 - http_cache = None + _msal_http_cache = None + + # Instance of ServicePrincipalStore. + # It follows singleton pattern so that _secret_file is read only once. + _service_principal_store_instance = None def __init__(self, authority, tenant_id=None, client_id=None, encrypt=False): """ @@ -57,35 +66,53 @@ def __init__(self, authority, tenant_id=None, client_id=None, encrypt=False): config_dir = get_config_dir() self._token_cache_file = os.path.join(config_dir, "msal_token_cache") self._secret_file = os.path.join(config_dir, "service_principal_entries") - self._http_cache_file = os.path.join(config_dir, "msal_http_cache") - - # Prepare HTTP cache. - # https://github.com/AzureAD/microsoft-authentication-library-for-python/pull/407 - # if not Identity.http_cache: - # Identity.http_cache = self._load_http_cache() + self._http_cache_file = os.path.join(config_dir, "msal_http_cache.bin") + # We make _msal_app_instance an instance attribute, instead of a class attribute, + # because MSAL apps can have different tenant IDs. self._msal_app_instance = None - # Store for Service principal credential persistence - self._msal_secret_store = ServicePrincipalStore(self._secret_file, self.encrypt) - self._msal_app_kwargs = { + + @property + def _msal_app_kwargs(self): + """kwargs for creating UserCredential or ServicePrincipalCredential. + MSAL token cache and HTTP cache are lazily created. + """ + if not Identity._msal_token_cache: + Identity._msal_token_cache = self._load_msal_token_cache() + + if not Identity._msal_http_cache: + Identity._msal_http_cache = self._load_msal_http_cache() + + return { "authority": self._msal_authority, - "token_cache": self._load_msal_cache() - # "http_cache": Identity.http_cache + "token_cache": Identity._msal_token_cache, + "http_cache": Identity._msal_http_cache } - def _load_msal_cache(self): + @property + def _msal_app(self): + """A PublicClientApplication instance for user login/logout. + The instance is lazily created. + """ + if not self._msal_app_instance: + self._msal_app_instance = PublicClientApplication(self.client_id, **self._msal_app_kwargs) + return self._msal_app_instance + + def _load_msal_token_cache(self): # Store for user token persistence cache = load_persisted_token_cache(self._token_cache_file, self.encrypt) return cache - def _load_http_cache(self): + def _load_msal_http_cache(self): import atexit import pickle + logger.debug("_load_msal_http_cache: %s", self._http_cache_file) try: with open(self._http_cache_file, 'rb') as f: - persisted_http_cache = pickle.load(f) # Take a snapshot - except: # pylint: disable=bare-except + persisted_http_cache = pickle.load(f) + except (pickle.UnpicklingError, FileNotFoundError) as ex: + logger.debug("Failed to load MSAL HTTP cache: %s", ex) persisted_http_cache = {} # Ignore a non-exist or corrupted http_cache atexit.register(lambda: pickle.dump( # When exit, flush it back to the file. @@ -95,45 +122,44 @@ def _load_http_cache(self): return persisted_http_cache - def _build_persistent_msal_app(self): - # Initialize _msal_app for login and logout - msal_app = PublicClientApplication(self.client_id, **self._msal_app_kwargs) - return msal_app - @property - def msal_app(self): - if not self._msal_app_instance: - self._msal_app_instance = self._build_persistent_msal_app() - return self._msal_app_instance + def _service_principal_store(self): + """A ServicePrincipalStore instance for service principal entries persistence. + The instance is lazily created. + """ + if not Identity._service_principal_store_instance: + store = load_secret_store(self._secret_file, self.encrypt) + Identity._service_principal_store_instance = ServicePrincipalStore(store) + return Identity._service_principal_store_instance def login_with_auth_code(self, scopes, **kwargs): # Emit a warning to inform that a browser is opened. # Only show the path part of the URL and hide the query string. logger.warning("The default web browser has been opened at %s. Please continue the login in the web browser. " "If no web browser is available or if the web browser fails to open, use device code flow " - "with `az login --use-device-code`.", self.msal_app.authority.authorization_endpoint) + "with `az login --use-device-code`.", self._msal_app.authority.authorization_endpoint) from .util import read_response_templates success_template, error_template = read_response_templates() # For AAD, use port 0 to let the system choose arbitrary unused ephemeral port to avoid port collision # on port 8400 from the old design. However, ADFS only allows port 8400. - result = self.msal_app.acquire_token_interactive( + result = self._msal_app.acquire_token_interactive( scopes, prompt='select_account', port=8400 if self._is_adfs else None, success_template=success_template, error_template=error_template, **kwargs) return check_result(result) def login_with_device_code(self, scopes, **kwargs): - flow = self.msal_app.initiate_device_flow(scopes, **kwargs) + flow = self._msal_app.initiate_device_flow(scopes, **kwargs) if "user_code" not in flow: raise ValueError( "Fail to create device flow. Err: %s" % json.dumps(flow, indent=4)) logger.warning(flow["message"]) - result = self.msal_app.acquire_token_by_device_flow(flow, **kwargs) # By default it will block + result = self._msal_app.acquire_token_by_device_flow(flow, **kwargs) # By default it will block return check_result(result) def login_with_username_password(self, username, password, scopes, **kwargs): - result = self.msal_app.acquire_token_by_username_password(username, password, scopes, **kwargs) + result = self._msal_app.acquire_token_by_username_password(username, password, scopes, **kwargs) return check_result(result) def login_with_service_principal(self, client_id, credential, scopes): @@ -149,7 +175,7 @@ def login_with_service_principal(self, client_id, credential, scopes): # Only persist the service principal after a successful login entry = sp_auth.get_entry_to_persist() - self._msal_secret_store.save_entry(entry) + self._service_principal_store.save_entry(entry) def login_with_managed_identity(self, scopes, identity_id=None): # pylint: disable=too-many-statements raise NotImplementedError @@ -158,9 +184,9 @@ def login_in_cloud_shell(self, scopes): raise NotImplementedError def logout_user(self, user): - accounts = self.msal_app.get_accounts(user) + accounts = self._msal_app.get_accounts(user) for account in accounts: - self.msal_app.remove_account(account) + self._msal_app.remove_account(account) def logout_all_users(self): for e in file_extensions.values(): @@ -168,27 +194,28 @@ def logout_all_users(self): def logout_service_principal(self, sp): # remove service principal secrets - self._msal_secret_store.remove_entry(sp) + self._service_principal_store.remove_entry(sp) def logout_all_service_principal(self): # remove service principal secrets - self._msal_secret_store.remove_all_entries() + for e in file_extensions.values(): + _try_remove(self._secret_file + e) 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): return UserCredential(self.client_id, username, **self._msal_app_kwargs) def get_service_principal_credential(self, client_id): - entry = self._msal_secret_store.load_entry(client_id, self.tenant_id) + entry = self._service_principal_store.load_entry(client_id, self.tenant_id) sp_auth = ServicePrincipalAuth(entry) return ServicePrincipalCredential(sp_auth, **self._msal_app_kwargs) def get_service_principal_entry(self, client_id): """This method is only used by --sdk-auth. DO NOT use it elsewhere.""" - return self._msal_secret_store.load_entry(client_id, self.tenant_id) + return self._service_principal_store.load_entry(client_id, self.tenant_id) def get_managed_identity_credential(self, client_id=None): raise NotImplementedError @@ -255,10 +282,8 @@ class ServicePrincipalStore: """Save secrets in MSAL custom secret store for Service Principal authentication. """ - def __init__(self, secret_file, encrypt): - from .persistence import load_secret_store - self._secret_store = load_secret_store(secret_file, encrypt) - self._secret_file = secret_file + def __init__(self, secret_store): + self._secret_store = secret_store self._entries = [] def load_entry(self, sp_id, tenant): @@ -305,10 +330,6 @@ def remove_entry(self, sp_id): if state_changed: self._save_persistence() - def remove_all_entries(self): - for e in file_extensions.values(): - _try_remove(self._secret_file + e) - def _save_persistence(self): self._secret_store.save(self._entries) diff --git a/src/azure-cli-core/azure/cli/core/auth/persistence.py b/src/azure-cli-core/azure/cli/core/auth/persistence.py index 560f7b4e78f..6ebec24883f 100644 --- a/src/azure-cli-core/azure/cli/core/auth/persistence.py +++ b/src/azure-cli-core/azure/cli/core/auth/persistence.py @@ -35,6 +35,7 @@ def load_secret_store(location, encrypt): def build_persistence(location, encrypt): """Build a suitable persistence instance based your current OS""" location += file_extensions[encrypt] + logger.debug("build_persistence: location=%r, encrypt=%r", location, encrypt) if encrypt: if sys.platform.startswith('win'): return FilePersistenceWithDataProtection(location) diff --git a/src/azure-cli-core/azure/cli/core/auth/tests/test_identity.py b/src/azure-cli-core/azure/cli/core/auth/tests/test_identity.py index 206fd28fd4a..6ca990cebd8 100644 --- a/src/azure-cli-core/azure/cli/core/auth/tests/test_identity.py +++ b/src/azure-cli-core/azure/cli/core/auth/tests/test_identity.py @@ -200,31 +200,25 @@ class TestMsalSecretStore(unittest.TestCase): 'client_secret': 'test_secret' } - @mock.patch('azure.cli.core.auth.persistence.load_secret_store') - def test_load_entry(self, load_secret_store_mock): + def test_load_entry(self): store = MemoryStore() - load_secret_store_mock.return_value = store - secret_store = ServicePrincipalStore(None, None) + secret_store = ServicePrincipalStore(store) store._content = [self.test_sp] entry = secret_store.load_entry("myapp", "mytenant") self.assertEqual(entry['client_secret'], "test_secret") - @mock.patch('azure.cli.core.auth.persistence.load_secret_store') - def test_save_entry(self, load_secret_store_mock): + def test_save_entry(self): store = MemoryStore() - load_secret_store_mock.return_value = store - secret_store = ServicePrincipalStore(None, None) + secret_store = ServicePrincipalStore(store) secret_store.save_entry(self.test_sp) assert store._content == [self.test_sp] - @mock.patch('azure.cli.core.auth.persistence.load_secret_store') - def test_save_entry_add_new(self, load_secret_store_mock): + def test_save_entry_add_new(self): store = MemoryStore() - load_secret_store_mock.return_value = store test_sp2 = { 'client_id': "myapp2", @@ -233,30 +227,26 @@ def test_save_entry_add_new(self, load_secret_store_mock): } store._content = [self.test_sp] - secret_store = ServicePrincipalStore(None, None) + secret_store = ServicePrincipalStore(store) secret_store.save_entry(test_sp2) assert store._content == [self.test_sp, test_sp2] - @mock.patch('azure.cli.core.auth.persistence.load_secret_store') - def test_save_entry_update_existing(self, load_secret_store_mock): + def test_save_entry_update_existing(self): store = MemoryStore() - load_secret_store_mock.return_value = store store._content = [self.test_sp] new_creds = self.test_sp.copy() new_creds['client_secret'] = 'test_secret' - secret_store = ServicePrincipalStore(None, None) + secret_store = ServicePrincipalStore(store) secret_store.save_entry(new_creds) assert store._content == [new_creds] - @mock.patch('azure.cli.core.auth.persistence.load_secret_store') - def test_remove_entry(self, load_secret_store_mock): + def test_remove_entry(self): store = MemoryStore() - load_secret_store_mock.return_value = store store._content = [self.test_sp] - secret_store = ServicePrincipalStore(None, None) + secret_store = ServicePrincipalStore(store) secret_store.remove_entry('myapp') assert store._content == [] diff --git a/src/azure-cli-core/azure/cli/core/tests/test_profile.py b/src/azure-cli-core/azure/cli/core/tests/test_profile.py index 6c654dbf58b..da8431d4fb9 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_profile.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_profile.py @@ -33,7 +33,7 @@ BEARER = 'Bearer' -class MockCredential: +class CredentialMock: def __init__(self, *args, **kwargs): super().__init__() @@ -46,6 +46,12 @@ def get_token(self, *scopes, **kwargs): return AccessToken(MOCK_ACCESS_TOKEN, MOCK_EXPIRES_ON) +# Used as the return_value of azure.cli.core.auth.identity.Identity.get_user_credential +# If we directly patch azure.cli.core.auth.msal_authentication.UserCredential with CredentialMock, +# get_user_credential will prepare MSAL token cache and HTTP cache which is time-consuming and unnecessary. +credential_mock = CredentialMock() + + class MSRestAzureAuthStub: def __init__(self, *args, **kwargs): self._token = { @@ -843,8 +849,8 @@ def test_get_current_account_user(self): self.assertEqual(user, self.user1) - @mock.patch('azure.cli.core.auth.identity.UserCredential', MockCredential) - def test_get_login_credentials(self): + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', return_value=credential_mock) + def test_get_login_credentials(self, get_user_credential_mock): cli = DummyCli() # setup storage_mock = {'subscriptions': None} @@ -858,6 +864,7 @@ def test_get_login_credentials(self): profile._set_subscriptions(consolidated) # action cred, subscription_id, _ = profile.get_login_credentials() + get_user_credential_mock.assert_called_with(self.user1) # verify self.assertEqual(subscription_id, test_subscription_id) @@ -866,8 +873,8 @@ def test_get_login_credentials(self): token = cred.get_token() self.assertEqual(token.token, MOCK_ACCESS_TOKEN) - @mock.patch('azure.cli.core.auth.identity.UserCredential', MockCredential) - def test_get_login_credentials_aux_subscriptions(self): + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', return_value=credential_mock) + def test_get_login_credentials_aux_subscriptions(self, get_user_credential_mock): cli = DummyCli() storage_mock = {'subscriptions': None} @@ -895,8 +902,8 @@ def test_get_login_credentials_aux_subscriptions(self): self.assertEqual(token.token, MOCK_ACCESS_TOKEN) self.assertEqual(aux_tokens[0].token, MOCK_ACCESS_TOKEN) - @mock.patch('azure.cli.core.auth.identity.UserCredential', MockCredential) - def test_get_login_credentials_aux_tenants(self): + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', return_value=credential_mock) + def test_get_login_credentials_aux_tenants(self, get_user_credential_mock): cli = DummyCli() storage_mock = {'subscriptions': None} @@ -1026,8 +1033,8 @@ def test_get_login_credentials_msi_user_assigned_with_res_id(self): self.assertTrue(cred.token_read_count) self.assertTrue(cred.msi_res_id, test_res_id) - @mock.patch('azure.cli.core.auth.identity.UserCredential', MockCredential) - def test_get_raw_token(self): + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', return_value=credential_mock) + def test_get_raw_token(self, get_user_credential_mock): cli = DummyCli() # setup storage_mock = {'subscriptions': None} @@ -1068,7 +1075,7 @@ def test_get_raw_token(self): @mock.patch('azure.cli.core.auth.identity.Identity.get_service_principal_credential') def test_get_raw_token_for_sp(self, get_service_principal_credential_mock): - get_service_principal_credential_mock.return_value = MockCredential() + get_service_principal_credential_mock.return_value = CredentialMock() cli = DummyCli() # setup storage_mock = {'subscriptions': None} diff --git a/src/azure-cli-core/setup.py b/src/azure-cli-core/setup.py index 664be8b97b9..75c9d79b412 100644 --- a/src/azure-cli-core/setup.py +++ b/src/azure-cli-core/setup.py @@ -52,7 +52,7 @@ 'jmespath', 'knack~=0.9.0', 'msal-extensions>=0.3.0,<0.4', - 'msal>=1.15.0,<2.0.0', + 'msal>=1.16.0,<2.0.0', 'paramiko>=2.0.8,<3.0.0', 'pkginfo>=1.5.0.1', 'PyJWT>=2.1.0', diff --git a/src/azure-cli-testsdk/azure/cli/testsdk/patches.py b/src/azure-cli-testsdk/azure/cli/testsdk/patches.py index 0d331179309..76abe36a4b1 100644 --- a/src/azure-cli-testsdk/azure/cli/testsdk/patches.py +++ b/src/azure-cli-testsdk/azure/cli/testsdk/patches.py @@ -61,19 +61,22 @@ def _handle_load_cached_subscription(*args, **kwargs): # pylint: disable=unused def patch_retrieve_token_for_user(unit_test): - class UserCredentialMock: + def get_user_credential_mock(*args, **kwargs): + class UserCredentialMock: - def __init__(self, *args, **kwargs): - pass + def __init__(self, *args, **kwargs): + super().__init__() - def get_token(*args, **kwargs): # pylint: disable=unused-argument - from azure.core.credentials import AccessToken - import time - fake_raw_token = 'top-secret-token-for-you' - now = int(time.time()) - return AccessToken(fake_raw_token, now + 3600) + def get_token(*args, **kwargs): # pylint: disable=unused-argument + from azure.core.credentials import AccessToken + import time + fake_raw_token = 'top-secret-token-for-you' + now = int(time.time()) + return AccessToken(fake_raw_token, now + 3600) - mock_in_unit_test(unit_test, 'azure.cli.core.auth.identity.UserCredential', UserCredentialMock) + return UserCredentialMock() + + mock_in_unit_test(unit_test, 'azure.cli.core.auth.identity.Identity.get_user_credential', get_user_credential_mock) def patch_long_run_operation_delay(unit_test): diff --git a/src/azure-cli/requirements.py3.Darwin.txt b/src/azure-cli/requirements.py3.Darwin.txt index 77ca75da0b2..72123634c0e 100644 --- a/src/azure-cli/requirements.py3.Darwin.txt +++ b/src/azure-cli/requirements.py3.Darwin.txt @@ -111,7 +111,7 @@ jsondiff==1.3.0 knack==0.9.0 MarkupSafe==1.1.1 msal-extensions==0.3.0 -msal==1.15.0 +msal==1.16.0 msrest==0.6.21 msrestazure==0.6.3 oauthlib==3.0.1 diff --git a/src/azure-cli/requirements.py3.Linux.txt b/src/azure-cli/requirements.py3.Linux.txt index 5b275f61fa1..8e5b7df159f 100644 --- a/src/azure-cli/requirements.py3.Linux.txt +++ b/src/azure-cli/requirements.py3.Linux.txt @@ -112,7 +112,7 @@ jsondiff==1.3.0 knack==0.9.0 MarkupSafe==1.1.1 msal-extensions==0.3.0 -msal==1.15.0 +msal==1.16.0 msrest==0.6.21 msrestazure==0.6.3 oauthlib==3.0.1 diff --git a/src/azure-cli/requirements.py3.windows.txt b/src/azure-cli/requirements.py3.windows.txt index 98ac4b2a417..e7afe9934b4 100644 --- a/src/azure-cli/requirements.py3.windows.txt +++ b/src/azure-cli/requirements.py3.windows.txt @@ -110,7 +110,7 @@ jsondiff==1.3.0 knack==0.9.0 MarkupSafe==1.1.1 msal-extensions==0.3.0 -msal==1.15.0 +msal==1.16.0 msrest==0.6.21 msrestazure==0.6.3 oauthlib==3.0.1 From 2d07b701a47b13a9ce1ec570d16ac39181a7b512 Mon Sep 17 00:00:00 2001 From: Jiashuo Li <4003950+jiasli@users.noreply.github.com> Date: Tue, 16 Nov 2021 17:03:14 +0800 Subject: [PATCH 2/2] Update identity.py --- src/azure-cli-core/azure/cli/core/auth/identity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/azure-cli-core/azure/cli/core/auth/identity.py b/src/azure-cli-core/azure/cli/core/auth/identity.py index 9d8f281e8b5..dd9575b4cb8 100644 --- a/src/azure-cli-core/azure/cli/core/auth/identity.py +++ b/src/azure-cli-core/azure/cli/core/auth/identity.py @@ -37,7 +37,7 @@ class Identity: # pylint: disable=too-many-instance-attributes _msal_token_cache = None # MSAL HTTP cache for MSAL's tenant discovery, retry-after error cache, etc. - # It *must* follow singleton pattern so that all MSAL app instances shares the same HTTP cache. + # It *must* follow singleton pattern so that all MSAL app instances share the same HTTP cache. # https://github.com/AzureAD/microsoft-authentication-library-for-python/pull/407 _msal_http_cache = None