{Core} Optimize performance for MSAL migration

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

@@ -10,8 +10,6 @@
 from enum import Enum
 
 from azure.cli.core._session import ACCOUNT
-from azure.cli.core.auth.identity import Identity, AZURE_CLI_CLIENT_ID
-from azure.cli.core.auth.util import resource_to_scopes
 from azure.cli.core.azclierror import AuthenticationError
 from azure.cli.core.cloud import get_active_cloud, set_cloud_subscription
 from azure.cli.core.util import in_cloud_console, can_launch_browser
@@ -121,12 +119,9 @@ def __init__(self, cli_ctx=None, storage=None):
         self.cli_ctx = cli_ctx or get_default_cli()
         self._storage = storage or ACCOUNT
         self._authority = self.cli_ctx.cloud.endpoints.active_directory
-        self._arm_scope = resource_to_scopes(self.cli_ctx.cloud.endpoints.active_directory_resource_id)
 
-        # Only enable token cache encryption for Windows (for now)
-        token_encryption_fallback = sys.platform.startswith('win32')
-        Identity.token_encryption = self.cli_ctx.config.getboolean('core', 'token_encryption',
-                                                                   fallback=token_encryption_fallback)
+        from .auth.util import resource_to_scopes
+        self._arm_scope = resource_to_scopes(self.cli_ctx.cloud.endpoints.active_directory_resource_id)
 
     # pylint: disable=too-many-branches,too-many-statements,too-many-locals
     def login(self,
@@ -136,7 +131,6 @@ def login(self,
               is_service_principal,
               tenant,
               scopes=None,
-              client_id=AZURE_CLI_CLIENT_ID,
               use_device_code=False,
               allow_no_subscriptions=False,
               use_cert_sn_issuer=None,
@@ -147,7 +141,7 @@ def login(self,
         if not scopes:
             scopes = self._arm_scope
 
-        identity = Identity(self._authority, tenant_id=tenant, client_id=client_id)
+        identity = _create_identity_instance(self.cli_ctx, self._authority, tenant_id=tenant)
 
         user_identity = None
         if interactive:
@@ -292,14 +286,14 @@ def logout(self, user_or_sp):
         subscriptions = [x for x in subscriptions if x not in result]
         self._storage[_SUBSCRIPTIONS] = subscriptions
 
-        identity = Identity(self._authority)
+        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] = []
 
-        identity = Identity(self._authority)
+        identity = _create_identity_instance(self.cli_ctx, self._authority)
         identity.logout_all_users()
         identity.logout_all_service_principal()
 
@@ -356,11 +350,12 @@ def get_login_credentials(self, resource=None, client_id=None, subscription_id=N
     def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=None):
         # Convert resource to scopes
         if resource and not scopes:
+            from .auth.util import resource_to_scopes
             scopes = resource_to_scopes(resource)
 
         # Use ARM as the default scopes
         if not scopes:
-            scopes = resource_to_scopes(self.cli_ctx.cloud.endpoints.active_directory_resource_id)
+            scopes = self._arm_scope
 
         if subscription and tenant:
             raise CLIError("Please specify only one of subscription and tenant, not both")
@@ -586,7 +581,7 @@ def _create_credential(self, account, tenant_id=None, client_id=None):
         user_type = account[_USER_ENTITY][_USER_TYPE]
         username_or_sp_id = account[_USER_ENTITY][_USER_NAME]
         tenant_id = tenant_id if tenant_id else account[_TENANT_ID]
-        identity = Identity(self._authority, tenant_id=tenant_id, client_id=client_id)
+        identity = _create_identity_instance(self.cli_ctx, self._authority, tenant_id=tenant_id, client_id=client_id)
 
         # User
         if user_type == _USER:
@@ -662,7 +657,7 @@ def get_sp_auth_info(self, subscription_id=None, name=None, password=None, cert_
             if user_type == _SERVICE_PRINCIPAL:
                 client_id = account[_USER_ENTITY][_USER_NAME]
                 result['clientId'] = client_id
-                identity = Identity(tenant_id=account[_TENANT_ID])
+                identity = _create_identity_instance(self.cli_ctx, self._authority, tenant_id=account[_TENANT_ID])
                 sp_entry = identity.get_service_principal_entry(client_id)
 
                 from .auth.msal_authentication import _CLIENT_SECRET, _CERTIFICATE
@@ -743,7 +738,7 @@ def __init__(self, cli_ctx):
         self.cli_ctx = cli_ctx
         self.secret = None
         self._arm_resource_id = cli_ctx.cloud.endpoints.active_directory_resource_id
-        self.authority = self.cli_ctx.cloud.endpoints.active_directory
+        self._authority = self.cli_ctx.cloud.endpoints.active_directory
         self.tenants = []
 
     def find_using_common_tenant(self, username, credential=None):
@@ -769,7 +764,7 @@ def find_using_common_tenant(self, username, credential=None):
 
             logger.info("Finding subscriptions under tenant %s", t.tenant_id_name)
 
-            identity = Identity(self.authority, tenant_id=tenant_id)
+            identity = _create_identity_instance(self.cli_ctx, self._authority, tenant_id=tenant_id)
 
             specific_tenant_credential = identity.get_user_credential(username)
 
@@ -863,3 +858,14 @@ def _transform_subscription_for_multiapi(s, s_dict):
             s_dict[_MANAGED_BY_TENANTS] = None
         else:
             s_dict[_MANAGED_BY_TENANTS] = [{_TENANT_ID: t.tenant_id} for t in s.managed_by_tenants]
+
+
+def _create_identity_instance(cli_ctx, *args, **kwargs):
+    """Lazily import and create Identity instance to avoid unnecessary imports."""
+    from .auth.identity import Identity
+
+    # Only enable encryption for Windows (for now).
+    fallback = sys.platform.startswith('win32')
+    encrypt = cli_ctx.config.getboolean('core', 'token_encryption', fallback=fallback)
+
+    return Identity(*args, encrypt=encrypt, **kwargs)

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

@@ -8,16 +8,16 @@
 import re
 
 from azure.cli.core._environment import get_config_dir
+from msal import PublicClientApplication
+
 from knack.log import get_logger
 from knack.util import CLIError
-
-from .msal_authentication import UserCredential, ServicePrincipalCredential
-from .util import check_result
-from .persistence import load_persisted_token_cache, file_extensions
-
 # Service principal entry properties
-from .msal_authentication import _CLIENT_ID, _TENANT, _CLIENT_SECRET, _CERTIFICATE, _CLIENT_ASSERTION,\
+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 .util import check_result
 
 AZURE_CLI_CLIENT_ID = '04b07795-8ddb-461a-bbee-02f9e1bf7b46'
 
@@ -31,26 +31,25 @@ class Identity:  # pylint: disable=too-many-instance-attributes
         - service principal
         - TODO: managed identity
     """
-    # Whether token and secrets should be encrypted. Change its value to turn on/off token encryption.
-    token_encryption = False
-
     # 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.
     # https://github.com/AzureAD/microsoft-authentication-library-for-python/pull/407
     http_cache = None
 
-    def __init__(self, authority, tenant_id=None, client_id=None):
+    def __init__(self, authority, tenant_id=None, client_id=None, encrypt=False):
         """
         :param authority: Authentication authority endpoint. For example,
             - AAD: https://login.microsoftonline.com
             - ADFS: https://adfs.redmond.azurestack.corp.microsoft.com/adfs
         :param tenant_id: Tenant GUID, like 00000000-0000-0000-0000-000000000000. If unspecified, default to
             'organizations'.
         :param client_id: Client ID of the CLI application.
+        :param encrypt:  Whether to encrypt token cache and service principal entries.
         """
         self.authority = authority
         self.tenant_id = tenant_id
         self.client_id = client_id or AZURE_CLI_CLIENT_ID
+        self.encrypt = encrypt
 
         # Build the authority in MSAL style
         self._msal_authority, self._is_adfs = _get_authority_url(authority, tenant_id)
@@ -67,7 +66,7 @@ def __init__(self, authority, tenant_id=None, client_id=None):
 
         self._msal_app_instance = None
         # Store for Service principal credential persistence
-        self._msal_secret_store = ServicePrincipalStore(self._secret_file, self.token_encryption)
+        self._msal_secret_store = ServicePrincipalStore(self._secret_file, self.encrypt)
         self._msal_app_kwargs = {
             "authority": self._msal_authority,
             "token_cache": self._load_msal_cache()
@@ -76,7 +75,7 @@ def __init__(self, authority, tenant_id=None, client_id=None):
 
     def _load_msal_cache(self):
         # Store for user token persistence
-        cache = load_persisted_token_cache(self._token_cache_file, self.token_encryption)
+        cache = load_persisted_token_cache(self._token_cache_file, self.encrypt)
         return cache
 
     def _load_http_cache(self):
@@ -98,7 +97,6 @@ def _load_http_cache(self):
 
     def _build_persistent_msal_app(self):
         # Initialize _msal_app for login and logout
-        from msal import PublicClientApplication
         msal_app = PublicClientApplication(self.client_id, **self._msal_app_kwargs)
         return msal_app
 

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

@@ -10,7 +10,6 @@
 SDK invocation.
 """
 
-from azure.core.credentials import AccessToken
 from knack.log import get_logger
 from knack.util import CLIError
 from msal import PublicClientApplication, ConfidentialClientApplication
@@ -103,4 +102,9 @@ def _build_sdk_access_token(token_entry):
     import time
     request_time = int(time.time())
 
+    # Importing azure.core.credentials.AccessToken is expensive.
+    # This can slow down commands that doesn't need azure.core, like `az account get-access-token`.
+    # So We define our own AccessToken.
+    from collections import namedtuple
+    AccessToken = namedtuple("AccessToken", ["token", "expires_on"])
     return AccessToken(token_entry["access_token"], request_time + token_entry["expires_in"])