From 4aef1a49c829f6172bf2ba433f7fd221c615c3d9 Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Fri, 14 Apr 2017 13:37:32 -0700 Subject: [PATCH 1/7] Performance: make subscriptionid, current user name available from the Profile without having to load adal and friends --- src/azure-cli-core/azure/cli/core/_profile.py | 49 +++++++++++-------- .../azure/cli/core/telemetry.py | 2 +- 2 files changed, 30 insertions(+), 21 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index aafb00f2583..1df1154ab7b 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -67,8 +67,8 @@ def _authentication_context_factory(authority, cache): CLOUD = get_active_cloud() logger.debug("Current active cloud '%s'", CLOUD.name) -logger.debug(pformat(vars(CLOUD.endpoints))) -logger.debug(pformat(vars(CLOUD.suffixes))) +logger.debug(vars(CLOUD.endpoints)) +logger.debug(vars(CLOUD.suffixes)) def get_authority_url(tenant=None): @@ -98,10 +98,10 @@ class CredentialType(Enum): # pylint: disable=too-few-public-methods class Profile(object): def __init__(self, storage=None, auth_ctx_factory=None): self._storage = storage or ACCOUNT - factory = auth_ctx_factory or _AUTH_CTX_FACTORY - self._creds_cache = CredsCache(factory) - self._subscription_finder = SubscriptionFinder(factory, self._creds_cache.adal_token_cache) + self.auth_ctx_factory = auth_ctx_factory or _AUTH_CTX_FACTORY + self._creds_cache = CredsCache(self.auth_ctx_factory) self._management_resource_uri = CLOUD.endpoints.management + self._subscription_finder_attr = None def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments interactive, @@ -113,17 +113,17 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments allow_debug_adal_connection() subscriptions = [] if interactive: - subscriptions = self._subscription_finder.find_through_interactive_flow( + subscriptions = self.subscription_finder.find_through_interactive_flow( tenant, self._management_resource_uri) else: if is_service_principal: if not tenant: raise CLIError('Please supply tenant using "--tenant"') sp_auth = ServicePrincipalAuth(password) - subscriptions = self._subscription_finder.find_from_service_principal_id( + subscriptions = self.subscription_finder.find_from_service_principal_id( username, sp_auth, tenant, self._management_resource_uri) else: - subscriptions = self._subscription_finder.find_from_user_account( + subscriptions = self.subscription_finder.find_from_user_account( username, password, tenant, self._management_resource_uri) if not subscriptions: @@ -135,13 +135,20 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments if self._creds_cache.adal_token_cache.has_state_changed: self._creds_cache.persist_cached_creds() - consolidated = Profile._normalize_properties(self._subscription_finder.user_id, + consolidated = Profile._normalize_properties(self.subscription_finder.user_id, subscriptions, is_service_principal) self._set_subscriptions(consolidated) # use deepcopy as we don't want to persist these changes to file. return deepcopy(consolidated) + @property + def subscription_finder(self): + if self._subscription_finder_attr is None: + self._subscription_finder_attr = SubscriptionFinder(self.auth_ctx_factory, + self._creds_cache.adal_token_cache) + return self._subscription_finder_attr + @staticmethod def _normalize_properties(user, subscriptions, is_service_principal): consolidated = [] @@ -259,6 +266,9 @@ def get_subscription(self, subscription=None): # take id or name raise CLIError("Please run 'az account set' to select active account.") return result[0] + def get_subscription_id(self): + return self.get_subscription()[_SUBSCRIPTION_ID] + def get_login_credentials(self, resource=CLOUD.endpoints.management, subscription_id=None): account = self.get_subscription(subscription_id) @@ -431,8 +441,7 @@ def __init__(self, auth_ctx_factory=None): os.path.join(get_config_dir(), 'accessTokens.json')) self._service_principal_creds = [] self._auth_ctx_factory = auth_ctx_factory or _AUTH_CTX_FACTORY - self.adal_token_cache = None - self._load_creds() + self._adal_token_cache_attr = None def persist_cached_creds(self): with os.fdopen(os.open(self._token_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o600), @@ -480,15 +489,15 @@ def retrieve_secret_of_service_principal(self, sp_id): cred = matched[0] return cred[_ACCESS_TOKEN] - def _load_creds(self): - import adal - if self.adal_token_cache is not None: - return self.adal_token_cache - all_entries = _load_tokens_from_file(self._token_file) - self._load_service_principal_creds(all_entries) - real_token = [x for x in all_entries if x not in self._service_principal_creds] - self.adal_token_cache = adal.TokenCache(json.dumps(real_token)) - return self.adal_token_cache + @property + def adal_token_cache(self): + if self._adal_token_cache_attr is None: + import adal + all_entries = _load_tokens_from_file(self._token_file) + self._load_service_principal_creds(all_entries) + real_token = [x for x in all_entries if x not in self._service_principal_creds] + self._adal_token_cache_attr = adal.TokenCache(json.dumps(real_token)) + return self._adal_token_cache_attr def save_service_principal_cred(self, sp_entry): matched = [x for x in self._service_principal_creds diff --git a/src/azure-cli-core/azure/cli/core/telemetry.py b/src/azure-cli-core/azure/cli/core/telemetry.py index c5cdc53d219..f5ae41f25b9 100644 --- a/src/azure-cli-core/azure/cli/core/telemetry.py +++ b/src/azure-cli-core/azure/cli/core/telemetry.py @@ -322,7 +322,7 @@ def _get_env_string(): @decorators.suppress_all_exceptions(fallback_return=None) def _get_azure_subscription_id(): - return _get_profile().get_login_credentials()[1] + return _get_profile().get_subscription_id() def _get_shell_type(): From d87a60e4c16c29d393a1a48937317808dda7b774 Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Fri, 14 Apr 2017 14:33:45 -0700 Subject: [PATCH 2/7] Make pylint happy (forgot to remove pprint import after I removed usage of pprint.pformat) --- src/azure-cli-core/azure/cli/core/_profile.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 1df1154ab7b..63424ea1b77 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -9,7 +9,6 @@ import errno import json import os.path -from pprint import pformat from copy import deepcopy from enum import Enum From af1e90ae87353fb329c4933959e7be62b0baeb61 Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Fri, 14 Apr 2017 18:15:04 -0700 Subject: [PATCH 3/7] Fix issues introduced by not loading the service principals cred cache (delay loaded a bit too much) --- src/azure-cli-core/azure/cli/core/_profile.py | 10 ++++++++-- src/azure-cli-core/tests/test_profile.py | 4 ++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 63424ea1b77..5aa12ee09ac 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -132,7 +132,7 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments self._creds_cache.save_service_principal_cred(sp_auth.get_entry_to_persist(username, tenant)) - if self._creds_cache.adal_token_cache.has_state_changed: + if self._creds_cache.load_adal_token_cache().has_state_changed: self._creds_cache.persist_cached_creds() consolidated = Profile._normalize_properties(self.subscription_finder.user_id, subscriptions, @@ -470,6 +470,7 @@ def retrieve_token_for_user(self, username, tenant, resource): return (token_entry[_TOKEN_ENTRY_TOKEN_TYPE], token_entry[_ACCESS_TOKEN]) def retrieve_token_for_service_principal(self, sp_id, resource): + self.load_adal_token_cache() matched = [x for x in self._service_principal_creds if sp_id == x[_SERVICE_PRINCIPAL_ID]] if not matched: raise CLIError("Please run 'az account set' to select active account.") @@ -482,6 +483,7 @@ def retrieve_token_for_service_principal(self, sp_id, resource): return (token_entry[_TOKEN_ENTRY_TOKEN_TYPE], token_entry[_ACCESS_TOKEN]) def retrieve_secret_of_service_principal(self, sp_id): + self.load_adal_token_cache() matched = [x for x in self._service_principal_creds if sp_id == x[_SERVICE_PRINCIPAL_ID]] if not matched: raise CLIError("No matched service principal found") @@ -490,6 +492,9 @@ def retrieve_secret_of_service_principal(self, sp_id): @property def adal_token_cache(self): + return self.load_adal_token_cache() + + def load_adal_token_cache(self): if self._adal_token_cache_attr is None: import adal all_entries = _load_tokens_from_file(self._token_file) @@ -499,6 +504,7 @@ def adal_token_cache(self): return self._adal_token_cache_attr def save_service_principal_cred(self, sp_entry): + self.load_adal_token_cache() matched = [x for x in self._service_principal_creds if sp_entry[_SERVICE_PRINCIPAL_ID] == x[_SERVICE_PRINCIPAL_ID] and sp_entry[_SERVICE_PRINCIPAL_TENANT] == x[_SERVICE_PRINCIPAL_TENANT]] @@ -526,7 +532,7 @@ def _load_service_principal_creds(self, creds): def remove_cached_creds(self, user_or_sp): state_changed = False # clear AAD tokens - tokens = self.adal_token_cache.find({_TOKEN_ENTRY_USER_ID: user_or_sp}) + tokens = self.load_adal_token_cache().find({_TOKEN_ENTRY_USER_ID: user_or_sp}) if tokens: state_changed = True self.adal_token_cache.remove(tokens) diff --git a/src/azure-cli-core/tests/test_profile.py b/src/azure-cli-core/tests/test_profile.py index a78211b840a..ed1ee79373e 100644 --- a/src/azure-cli-core/tests/test_profile.py +++ b/src/azure-cli-core/tests/test_profile.py @@ -310,7 +310,6 @@ def test_get_login_credentials(self, mock_get_token, mock_read_cred_file): token_type, token = cred._token_retriever() self.assertEqual(token, self.raw_token1) self.assertEqual(some_token_type, token_type) - self.assertEqual(mock_read_cred_file.call_count, 1) mock_get_token.assert_called_once_with(mock.ANY, self.user1, self.tenant_id, 'https://management.core.windows.net/') self.assertEqual(mock_get_token.call_count, 1) @@ -535,7 +534,7 @@ def test_credscache_load_tokens_and_sp_creds_with_secret(self, mock_read_file): creds_cache = CredsCache() # assert - token_entries = [entry for _, entry in creds_cache.adal_token_cache.read_items()] + token_entries = [entry for _, entry in creds_cache.load_adal_token_cache().read_items()] self.assertEqual(token_entries, [self.token_entry1]) self.assertEqual(creds_cache._service_principal_creds, [test_sp]) @@ -550,6 +549,7 @@ def test_credscache_load_tokens_and_sp_creds_with_cert(self, mock_read_file): # action creds_cache = CredsCache() + creds_cache.load_adal_token_cache() # assert self.assertEqual(creds_cache._service_principal_creds, [test_sp]) From e8b993a2c6844d2a063368c8f0183c5c00a364b7 Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Fri, 14 Apr 2017 18:56:33 -0700 Subject: [PATCH 4/7] Update to use adal_token_cache property instead of the generic load method (cleaner code) --- src/azure-cli-core/azure/cli/core/_profile.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 5aa12ee09ac..56496940979 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -132,7 +132,7 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments self._creds_cache.save_service_principal_cred(sp_auth.get_entry_to_persist(username, tenant)) - if self._creds_cache.load_adal_token_cache().has_state_changed: + if self._creds_cache.adal_token_cache.has_state_changed: self._creds_cache.persist_cached_creds() consolidated = Profile._normalize_properties(self.subscription_finder.user_id, subscriptions, @@ -532,7 +532,7 @@ def _load_service_principal_creds(self, creds): def remove_cached_creds(self, user_or_sp): state_changed = False # clear AAD tokens - tokens = self.load_adal_token_cache().find({_TOKEN_ENTRY_USER_ID: user_or_sp}) + tokens = self.adal_token_cache.find({_TOKEN_ENTRY_USER_ID: user_or_sp}) if tokens: state_changed = True self.adal_token_cache.remove(tokens) From 7337dd9d27b778d5bd896084cc9dc5771c71ea21 Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Fri, 14 Apr 2017 19:01:42 -0700 Subject: [PATCH 5/7] Fix broken test (renamed attribute name in profile wasn't updated in the test) --- src/azure-cli-core/tests/test_profile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/azure-cli-core/tests/test_profile.py b/src/azure-cli-core/tests/test_profile.py index ed1ee79373e..77bbebf4489 100644 --- a/src/azure-cli-core/tests/test_profile.py +++ b/src/azure-cli-core/tests/test_profile.py @@ -238,7 +238,7 @@ def test_get_expanded_subscription_info_for_logged_in_service_principal(self, storage_mock = {'subscriptions': []} profile = Profile(storage_mock) profile._management_resource_uri = 'https://management.core.windows.net/' - profile._subscription_finder = finder + profile._subscription_finder_attr = finder profile.find_subscriptions_on_login(False, '1234', 'my-secret', True, self.tenant_id) # action extended_info = profile.get_expanded_subscription_info() From dce8f814737cc56a9b58d6443b6f451f576c5a1d Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Mon, 17 Apr 2017 11:05:11 -0700 Subject: [PATCH 6/7] Fix broken merge from upstream (renamed attribute) --- src/azure-cli-core/azure/cli/core/_profile.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 3799c12f4d6..3910bbbe582 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -111,17 +111,17 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments allow_debug_adal_connection() subscriptions = [] if interactive: - subscriptions = self._subscription_finder.find_through_interactive_flow( + subscriptions = self.subscription_finder.find_through_interactive_flow( tenant, self._ad_resource_uri) else: if is_service_principal: if not tenant: raise CLIError('Please supply tenant using "--tenant"') sp_auth = ServicePrincipalAuth(password) - subscriptions = self._subscription_finder.find_from_service_principal_id( + subscriptions = self.subscription_finder.find_from_service_principal_id( username, sp_auth, tenant, self._ad_resource_uri) else: - subscriptions = self._subscription_finder.find_from_user_account( + subscriptions = self.subscription_finder.find_from_user_account( username, password, tenant, self._ad_resource_uri) if not subscriptions: From 7fcebae01f2c64d7117bb4cd047c8c73ffa97777 Mon Sep 17 00:00:00 2001 From: Johan Stenberg Date: Tue, 18 Apr 2017 13:20:37 -0700 Subject: [PATCH 7/7] Address code review feedback (remove subscription_finder property) --- src/azure-cli-core/azure/cli/core/_profile.py | 23 ++++++++----------- src/azure-cli-core/tests/test_profile.py | 8 +++++-- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 3910bbbe582..e0d45196a51 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -98,7 +98,6 @@ def __init__(self, storage=None, auth_ctx_factory=None): self.auth_ctx_factory = auth_ctx_factory or _AUTH_CTX_FACTORY self._creds_cache = CredsCache(self.auth_ctx_factory) self._management_resource_uri = CLOUD.endpoints.management - self._subscription_finder_attr = None self._ad_resource_uri = CLOUD.endpoints.active_directory_resource_id def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments @@ -106,22 +105,27 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments username, password, is_service_principal, - tenant): + tenant, + subscription_finder=None): from azure.cli.core._debug import allow_debug_adal_connection allow_debug_adal_connection() subscriptions = [] + + if not subscription_finder: + subscription_finder = SubscriptionFinder(self.auth_ctx_factory, + self._creds_cache.adal_token_cache) if interactive: - subscriptions = self.subscription_finder.find_through_interactive_flow( + subscriptions = subscription_finder.find_through_interactive_flow( tenant, self._ad_resource_uri) else: if is_service_principal: if not tenant: raise CLIError('Please supply tenant using "--tenant"') sp_auth = ServicePrincipalAuth(password) - subscriptions = self.subscription_finder.find_from_service_principal_id( + subscriptions = subscription_finder.find_from_service_principal_id( username, sp_auth, tenant, self._ad_resource_uri) else: - subscriptions = self.subscription_finder.find_from_user_account( + subscriptions = subscription_finder.find_from_user_account( username, password, tenant, self._ad_resource_uri) if not subscriptions: @@ -133,20 +137,13 @@ def find_subscriptions_on_login(self, # pylint: disable=too-many-arguments if self._creds_cache.adal_token_cache.has_state_changed: self._creds_cache.persist_cached_creds() - consolidated = Profile._normalize_properties(self.subscription_finder.user_id, + consolidated = Profile._normalize_properties(subscription_finder.user_id, subscriptions, is_service_principal) self._set_subscriptions(consolidated) # use deepcopy as we don't want to persist these changes to file. return deepcopy(consolidated) - @property - def subscription_finder(self): - if self._subscription_finder_attr is None: - self._subscription_finder_attr = SubscriptionFinder(self.auth_ctx_factory, - self._creds_cache.adal_token_cache) - return self._subscription_finder_attr - @staticmethod def _normalize_properties(user, subscriptions, is_service_principal): consolidated = [] diff --git a/src/azure-cli-core/tests/test_profile.py b/src/azure-cli-core/tests/test_profile.py index 0c7769495e6..e344826e285 100644 --- a/src/azure-cli-core/tests/test_profile.py +++ b/src/azure-cli-core/tests/test_profile.py @@ -238,8 +238,12 @@ def test_get_expanded_subscription_info_for_logged_in_service_principal(self, storage_mock = {'subscriptions': []} profile = Profile(storage_mock) profile._management_resource_uri = 'https://management.core.windows.net/' - profile._subscription_finder_attr = finder - profile.find_subscriptions_on_login(False, '1234', 'my-secret', True, self.tenant_id) + profile.find_subscriptions_on_login(False, + '1234', + 'my-secret', + True, + self.tenant_id, + finder) # action extended_info = profile.get_expanded_subscription_info() # assert