update

Author: jiasli

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

@@ -57,6 +57,10 @@ def _prepare_msal_kwargs(options=None):
     # 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. 'data' support is a CLI-specific extension.
+        # SDK doesn't support 'data': https://github.com/Azure/azure-sdk-for-python/pull/16397
+        if 'data' in options:
+            msal_kwargs['data'] = options['data']
         # For CAE
         if 'claims' in options:
             msal_kwargs['claims_challenge'] = options['claims']

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

@@ -43,23 +43,16 @@ def __init__(self, client_id, username, **kwargs):
 
         self._account = accounts[0]
 
-    def acquire_token(self, scopes, claims_challenge=None, data=None, **kwargs):
+    def acquire_token(self, scopes, claims_challenge=None, **kwargs):
         # scopes must be a list.
         # For acquiring SSH certificate, scopes is ['https://pas.windows.net/CheckMyAccess/Linux/.default']
-        # data is only used for acquiring VM SSH certificate. DO NOT use it for other purposes.
         # kwargs is already sanitized by CredentialAdaptor, so it can be safely passed to MSAL
-        logger.debug("UserCredential.acquire_token: scopes=%r, claims_challenge=%r, data=%r, kwargs=%r",
-                     scopes, claims_challenge, data, kwargs)
+        logger.debug("UserCredential.acquire_token: scopes=%r, claims_challenge=%r, kwargs=%r",
+                     scopes, claims_challenge, kwargs)
 
         if claims_challenge:
             logger.warning('Acquiring new access token silently for tenant %s with claims challenge: %s',
                            self._msal_app.authority.tenant, claims_challenge)
-
-        # Only pass data to MSAL if it is set. Passing data=None will cause failure in MSAL:
-        #   AttributeError: 'NoneType' object has no attribute 'get'
-        if data is not None:
-            kwargs['data'] = data
-
         result = self._msal_app.acquire_token_silent_with_error(
             scopes, self._account, claims_challenge=claims_challenge, **kwargs)
 
@@ -112,13 +105,8 @@ def __init__(self, client_id, client_credential, **kwargs):
         """
         self._msal_app = ConfidentialClientApplication(client_id, client_credential=client_credential, **kwargs)
 
-    def acquire_token(self, scopes, data=None, **kwargs):
-        logger.debug("ServicePrincipalCredential.acquire_token: scopes=%r, data=%r, kwargs=%r",
-                     scopes, data, kwargs)
-
-        if data is not None:
-            kwargs['data'] = data
-
+    def acquire_token(self, scopes, **kwargs):
+        logger.debug("ServicePrincipalCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)
         result = self._msal_app.acquire_token_for_client(scopes, **kwargs)
         check_result(result)
         return result
@@ -138,13 +126,8 @@ def __init__(self):
             #   token_cache=...
         )
 
-    def acquire_token(self, scopes, data=None, **kwargs):
-        logger.debug("CloudShellCredential.acquire_token: scopes=%r, data=%r, kwargs=%r",
-                     scopes, data, kwargs)
-
-        if data is not None:
-            kwargs['data'] = data
-
+    def acquire_token(self, scopes, **kwargs):
+        logger.debug("CloudShellCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)
         result = self._msal_app.acquire_token_interactive(scopes, prompt="none", **kwargs)
         check_result(result, scopes=scopes)
         return result
@@ -164,13 +147,8 @@ def __init__(self, client_id=None, resource_id=None, object_id=None):
             managed_identity = SystemAssignedManagedIdentity()
         self._msal_client = ManagedIdentityClient(managed_identity, http_client=requests.Session())
 
-    def acquire_token(self, scopes, data=None, **kwargs):
-        logger.debug("ManagedIdentityCredential.acquire_token: scopes=%r, data=%r, kwargs=%r",
-                     scopes, data, kwargs)
-
-        if data is not None:
-            from azure.cli.core.azclierror import AuthenticationError
-            raise AuthenticationError("VM SSH currently doesn't support managed identity.")
+    def acquire_token(self, scopes, **kwargs):
+        logger.debug("ManagedIdentityCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)
 
         from .util import scopes_to_resource
         result = self._msal_client.acquire_token_for_client(resource=scopes_to_resource(scopes))

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

@@ -11,6 +11,11 @@
 
 
 MOCK_ACCESS_TOKEN = "mock_access_token"
+MOCK_DATA = {
+    'key_id': 'test',
+    'req_cnf': 'test',
+    'token_type': 'ssh-cert'
+}
 MOCK_CLAIMS = {"test_claims": "value2"}
 
 class MsalCredentialStub:
@@ -39,7 +44,7 @@ def _now_timestamp_mock():
 
 class TestCredentialAdaptor(unittest.TestCase):
 
-    @mock.patch('azure.cli.core.auth.util.now_timestamp', new=_now_timestamp_mock)
+    @mock.patch('azure.cli.core.auth.util._now_timestamp', new=_now_timestamp_mock)
     def test_get_token(self):
         msal_cred = MsalCredentialStub()
         sdk_cred = CredentialAdaptor(msal_cred)
@@ -51,11 +56,15 @@ def test_get_token(self):
         assert access_token.token == MOCK_ACCESS_TOKEN
         assert access_token.expires_on == 1630920323
 
+        # Note that SDK doesn't support 'data'. This is a CLI-specific extension.
+        sdk_cred.get_token('https://management.core.windows.net//.default', data=MOCK_DATA)
+        assert msal_cred.acquire_token_kwargs['data'] == MOCK_DATA
+
         sdk_cred.get_token('https://management.core.windows.net//.default', claims=MOCK_CLAIMS)
         assert msal_cred.acquire_token_claims_challenge == MOCK_CLAIMS
 
 
-    @mock.patch('azure.cli.core.auth.util.now_timestamp', new=_now_timestamp_mock)
+    @mock.patch('azure.cli.core.auth.util._now_timestamp', new=_now_timestamp_mock)
     def test_get_token_info(self):
         msal_cred = MsalCredentialStub()
         sdk_cred = CredentialAdaptor(msal_cred)
@@ -69,6 +78,10 @@ def test_get_token_info(self):
 
         assert msal_cred.acquire_token_scopes == ['https://management.core.windows.net//.default']
 
+        # Note that SDK doesn't support 'data'. If 'data' were supported, it should be tested with:
+        sdk_cred.get_token_info('https://management.core.windows.net//.default', options={'data': MOCK_DATA})
+        assert msal_cred.acquire_token_kwargs['data'] == MOCK_DATA
+
         sdk_cred.get_token_info('https://management.core.windows.net//.default', options={'claims': MOCK_CLAIMS})
         assert msal_cred.acquire_token_claims_challenge == MOCK_CLAIMS