_generate_msal_kwargs

Author: jiasli

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

@@ -27,25 +27,15 @@ def get_token(self, *scopes, **kwargs):
         """Get an access token from the main credential."""
         logger.debug("CredentialAdaptor.get_token: scopes=%r, kwargs=%r", scopes, kwargs)
 
-        # Discard unsupported kwargs: tenant_id, enable_cae
-        filtered_kwargs = {}
-        if 'data' in kwargs:
-            filtered_kwargs['data'] = kwargs['data']
-
-        return build_sdk_access_token(self._credential.acquire_token(list(scopes), **filtered_kwargs))
+        msal_kwargs = _generate_msal_kwargs(kwargs)
+        return build_sdk_access_token(self._credential.acquire_token(list(scopes), **msal_kwargs))
 
 
     def get_token_info(self, *scopes, options=None):
         """Get an access token from the main credential."""
         logger.debug("CredentialAdaptor.get_token_info: scopes=%r, options=%r", scopes, options)
 
-        # Discard unsupported kwargs: tenant_id, enable_cae
-        msal_kwargs = {}
-        if 'claims' in options:
-            msal_kwargs['claims'] = options['claims']
-        if 'data' in options:
-            msal_kwargs['data'] = options['data']
-
+        msal_kwargs = _generate_msal_kwargs(options)
         return _build_sdk_access_token_info(self._credential.acquire_token(list(scopes), **msal_kwargs))
 
 
@@ -58,6 +48,20 @@ def get_auxiliary_tokens(self, *scopes, **kwargs):
         return None
 
 
+def _generate_msal_kwargs(options=None):
+    # Preserve supported options and discard unsupported options (tenant_id, enable_cae).
+    # Both get_token's kwargs and get_token_info's options are accepted as their schema is the same (at least for now).
+    msal_kwargs = {}
+    if options:
+        # For VM SSH
+        if 'data' in options:
+            msal_kwargs['data'] = options['data']
+        # For CAE
+        if 'claims' in options:
+            msal_kwargs['claims'] = options['claims']
+    return msal_kwargs
+
+
 def _build_sdk_access_token_info(token_entry):
     # MSAL token entry sample:
     # {