Add mtls_proof_of_possession and with_attestation_support APIs to acquire_token_for_client

Co-authored-by: gladjohn <90415114+gladjohn@users.noreply.github.com>

Author: Copilot

msal/managed_identity.py

@@ -281,6 +281,8 @@ def acquire_token_for_client(
         *,
         resource: str,  # If/when we support scope, resource will become optional
         claims_challenge: Optional[str] = None,
+        mtls_proof_of_possession: bool = False,
+        with_attestation_support: bool = False,
     ):
         """Acquire token for the managed identity.
 
@@ -300,6 +302,22 @@ def acquire_token_for_client(
             even if the app developer did not opt in for the "CP1" client capability.
             Upon receiving a `claims_challenge`, MSAL will attempt to acquire a new token.
 
+        :param bool mtls_proof_of_possession: (optional)
+            When True, use the MSI v2 (mTLS Proof-of-Possession) flow to acquire an
+            ``mtls_pop`` token bound to a short-lived mTLS certificate issued by the
+            IMDS ``/issuecredential`` endpoint.
+            Without this flag the legacy IMDS v1 flow is used.
+            Defaults to False.
+
+            This takes precedence over the ``msi_v2_enabled`` constructor parameter.
+
+        :param bool with_attestation_support: (optional)
+            When True (and ``mtls_proof_of_possession`` is also True), attempt
+            KeyGuard / platform attestation before credential issuance.
+            On Windows this leverages ``AttestationClientLib.dll`` when available;
+            on other platforms the parameter is silently ignored.
+            Defaults to False.
+
         .. note::
 
             Known issue: When an Azure VM has only one user-assigned managed identity,
@@ -349,11 +367,15 @@ def acquire_token_for_client(
                 return access_token_from_cache  # It is still good as new
         try:
             result = None
-            if self._msi_v2_enabled:
+            # Per-call mtls_proof_of_possession takes precedence over the constructor
+            # default (msi_v2_enabled / MSAL_ENABLE_MSI_V2 env var).
+            use_msi_v2 = mtls_proof_of_possession or self._msi_v2_enabled
+            if use_msi_v2:
                 try:
                     from .msi_v2 import obtain_token as _obtain_token_v2
                     result = _obtain_token_v2(
-                        self._http_client, self._managed_identity, resource)
+                        self._http_client, self._managed_identity, resource,
+                        attestation_enabled=with_attestation_support)
                     logger.debug("MSI v2 token acquisition succeeded")
                 except MsiV2Error as exc:
                     logger.warning(

msal/msi_v2.py

@@ -282,12 +282,16 @@ def obtain_token(
     http_client,
     managed_identity,
     resource: str,
+    attestation_enabled: bool = False,
 ) -> Dict[str, Any]:
     """Acquire a token using the MSI v2 (mTLS PoP) flow.
 
     :param http_client: HTTP client for IMDS requests.
     :param managed_identity: ManagedIdentity configuration dict.
     :param resource: Resource URL for token acquisition.
+    :param attestation_enabled: When True, attempt KeyGuard / platform attestation
+        before issuing credentials (Windows only; silently skipped on other platforms).
+        Defaults to False.
     :returns: OAuth2 token response dict with access_token on success,
               or error dict on failure.
     :raises MsiV2Error: If the flow fails at a non-recoverable step.
@@ -311,9 +315,9 @@ def obtain_token(
     # 3. Build PKCS#10 CSR with cuId OID extension
     csr_der = _build_csr(private_key, client_id, cu_id)
 
-    # 4. Attempt attestation (Windows only; falls back to None on other platforms)
+    # 4. Attempt attestation only when explicitly requested by the caller
     attestation_jwt = None
-    if attestation_endpoint:
+    if attestation_enabled and attestation_endpoint:
         try:
             from .msi_v2_attestation import get_attestation_jwt
             attestation_jwt = get_attestation_jwt(

sample/msi_v2_sample.py

@@ -13,10 +13,6 @@
 - Set RESOURCE environment variable to the target resource URL, e.g.
     export RESOURCE=https://management.azure.com/
 
-To enable MSI v2 (required):
-    export MSAL_ENABLE_MSI_V2=true
-  or pass msi_v2_enabled=True to ManagedIdentityClient.
-
 Usage:
     python msi_v2_sample.py
 """
@@ -42,13 +38,16 @@
     msal.SystemAssignedManagedIdentity(),
     http_client=requests.Session(),
     token_cache=global_token_cache,
-    msi_v2_enabled=True,  # Enable MSI v2 (mTLS PoP) flow
 )
 
 
 def acquire_and_use_token():
     """Acquire an mtls_pop token via MSI v2 and optionally call an API."""
-    result = client.acquire_token_for_client(resource=RESOURCE)
+    result = client.acquire_token_for_client(
+        resource=RESOURCE,
+        mtls_proof_of_possession=True,  # Use MSI v2 (mTLS PoP) flow
+        with_attestation_support=True,  # Enable KeyGuard attestation (Windows)
+    )
 
     if "access_token" in result:
         print("Token acquired successfully")

tests/test_msi_v2.py

@@ -435,7 +435,7 @@ def test_client_msi_v2_disabled_via_env_var(self):
 class TestMsiV2TokenAcquisitionIntegration(unittest.TestCase):
     """Integration tests for MSI v2 token acquisition flow with mocked IMDS."""
 
-    def _make_client(self, msi_v2_enabled=True):
+    def _make_client(self, msi_v2_enabled=False):
         import requests
         return msal.ManagedIdentityClient(
             msal.SystemAssignedManagedIdentity(),
@@ -466,7 +466,48 @@ def _make_mock_responses(self, client_id, cu_id, cert_pem, token_endpoint,
 
     @patch("msal.msi_v2._acquire_token_via_mtls")
     def test_msi_v2_happy_path(self, mock_mtls):
-        """MSI v2 succeeds end-to-end (mTLS call is mocked)."""
+        """MSI v2 succeeds end-to-end via mtls_proof_of_possession=True."""
+        import requests
+
+        key = _generate_rsa_key()
+        cert_pem = _make_self_signed_cert(key, "test-client-id")
+        access_token = "MSI_V2_ACCESS_TOKEN"
+        expires_in = 3600
+        token_endpoint = "https://login.microsoftonline.com/tenant/oauth2/token"
+
+        platform_metadata, credential, token_response = self._make_mock_responses(
+            "test-client-id", "test-cu-id", cert_pem, token_endpoint,
+            access_token, expires_in)
+
+        mock_mtls.return_value = token_response
+
+        client = self._make_client()
+
+        def _mock_get(url, **kwargs):
+            if "getplatformmetadata" in url:
+                return MinimalResponse(
+                    status_code=200, text=json.dumps(platform_metadata))
+            raise ValueError("Unexpected GET: {}".format(url))
+
+        def _mock_post(url, **kwargs):
+            if "issuecredential" in url:
+                return MinimalResponse(
+                    status_code=200, text=json.dumps(credential))
+            raise ValueError("Unexpected POST: {}".format(url))
+
+        with patch.object(client._http_client, "get", side_effect=_mock_get), \
+             patch.object(client._http_client, "post", side_effect=_mock_post):
+            result = client.acquire_token_for_client(
+                resource="https://management.azure.com/",
+                mtls_proof_of_possession=True)
+
+        self.assertEqual(result["access_token"], access_token)
+        self.assertEqual(result["token_type"], "mtls_pop")
+        self.assertEqual(result["token_source"], "identity_provider")
+
+    @patch("msal.msi_v2._acquire_token_via_mtls")
+    def test_msi_v2_happy_path_via_constructor_flag(self, mock_mtls):
+        """MSI v2 also works when enabled via the msi_v2_enabled constructor param."""
         import requests
 
         key = _generate_rsa_key()
@@ -497,18 +538,18 @@ def _mock_post(url, **kwargs):
 
         with patch.object(client._http_client, "get", side_effect=_mock_get), \
              patch.object(client._http_client, "post", side_effect=_mock_post):
+            # No mtls_proof_of_possession kwarg; relies on constructor flag
             result = client.acquire_token_for_client(
                 resource="https://management.azure.com/")
 
         self.assertEqual(result["access_token"], access_token)
         self.assertEqual(result["token_type"], "mtls_pop")
-        self.assertEqual(result["token_source"], "identity_provider")
 
     @patch("msal.msi_v2._acquire_token_via_mtls")
     def test_msi_v2_fallback_to_v1_on_metadata_failure(self, mock_mtls):
         """MSI v2 falls back to MSI v1 if IMDS metadata call fails."""
         import requests
-        client = self._make_client(msi_v2_enabled=True)
+        client = self._make_client()
 
         def _mock_get(url, **kwargs):
             if "getplatformmetadata" in url:
@@ -523,18 +564,19 @@ def _mock_get(url, **kwargs):
             raise ValueError("Unexpected GET: {}".format(url))
 
         with patch.object(client._http_client, "get", side_effect=_mock_get):
-            result = client.acquire_token_for_client(resource="R")
+            result = client.acquire_token_for_client(
+                resource="R", mtls_proof_of_possession=True)
 
         # Should have fallen back to MSI v1
         self.assertEqual(result["access_token"], "V1_TOKEN")
         mock_mtls.assert_not_called()
 
     @patch("msal.msi_v2._acquire_token_via_mtls")
-    def test_msi_v2_not_attempted_when_disabled(self, mock_mtls):
-        """MSI v2 is not attempted when msi_v2_enabled=False."""
+    def test_msi_v2_not_attempted_when_not_requested(self, mock_mtls):
+        """MSI v2 is not attempted when mtls_proof_of_possession=False (default)."""
         import requests
 
-        client = self._make_client(msi_v2_enabled=False)
+        client = self._make_client()
 
         def _mock_get(url, **kwargs):
             return MinimalResponse(status_code=200, text=json.dumps({
@@ -544,6 +586,7 @@ def _mock_get(url, **kwargs):
             }))
 
         with patch.object(client._http_client, "get", side_effect=_mock_get):
+            # No mtls_proof_of_possession — uses v1 by default
             result = client.acquire_token_for_client(resource="R")
 
         mock_mtls.assert_not_called()
@@ -553,7 +596,7 @@ def _mock_get(url, **kwargs):
     def test_msi_v2_fallback_on_unexpected_error(self, mock_mtls):
         """MSI v2 falls back to MSI v1 on unexpected errors."""
         import requests
-        client = self._make_client(msi_v2_enabled=True)
+        client = self._make_client()
 
         platform_metadata = {
             "clientId": "client-id",
@@ -583,12 +626,85 @@ def _mock_post(url, **kwargs):
 
         with patch.object(client._http_client, "get", side_effect=_mock_get), \
              patch.object(client._http_client, "post", side_effect=_mock_post):
-            result = client.acquire_token_for_client(resource="R")
+            result = client.acquire_token_for_client(
+                resource="R", mtls_proof_of_possession=True)
 
         # Should fall back to MSI v1
         self.assertEqual(result["access_token"], "V1_FALLBACK")
         mock_mtls.assert_not_called()
 
+    @patch("msal.msi_v2_attestation.get_attestation_jwt")
+    @patch("msal.msi_v2._acquire_token_via_mtls")
+    def test_with_attestation_support_triggers_attestation(
+        self, mock_mtls, mock_attest
+    ):
+        """with_attestation_support=True calls attestation; False skips it."""
+        import requests
+
+        key = _generate_rsa_key()
+        cert_pem = _make_self_signed_cert(key, "test-client-id")
+        token_endpoint = "https://login.microsoftonline.com/tenant/oauth2/token"
+        access_token = "MSI_V2_ATTEST_TOKEN"
+        expires_in = 3600
+
+        platform_metadata = {
+            "clientId": "test-client-id",
+            "tenantId": "tenant-id",
+            "cuId": "test-cu-id",
+            "attestationEndpoint": "https://attest.example.com",
+        }
+        credential = {
+            "certificate": cert_pem,
+            "tokenEndpoint": token_endpoint,
+        }
+        token_response = {
+            "access_token": access_token,
+            "expires_in": str(expires_in),
+            "token_type": "mtls_pop",
+        }
+
+        mock_attest.return_value = "fake.attestation.jwt"
+        mock_mtls.return_value = token_response
+
+        client = self._make_client()
+
+        def _mock_get(url, **kwargs):
+            if "getplatformmetadata" in url:
+                return MinimalResponse(
+                    status_code=200, text=json.dumps(platform_metadata))
+            raise ValueError("Unexpected GET: {}".format(url))
+
+        def _mock_post(url, **kwargs):
+            if "issuecredential" in url:
+                return MinimalResponse(
+                    status_code=200, text=json.dumps(credential))
+            raise ValueError("Unexpected POST: {}".format(url))
+
+        # --- with_attestation_support=True: attestation should be called ---
+        with patch.object(client._http_client, "get", side_effect=_mock_get), \
+             patch.object(client._http_client, "post", side_effect=_mock_post):
+            result = client.acquire_token_for_client(
+                resource="https://management.azure.com/",
+                mtls_proof_of_possession=True,
+                with_attestation_support=True,
+            )
+        mock_attest.assert_called_once()
+        self.assertEqual(result["access_token"], access_token)
+
+        mock_attest.reset_mock()
+        mock_mtls.reset_mock()
+
+        # --- with_attestation_support=False (default): attestation NOT called ---
+        with patch.object(client._http_client, "get", side_effect=_mock_get), \
+             patch.object(client._http_client, "post", side_effect=_mock_post):
+            result = client.acquire_token_for_client(
+                resource="https://management.azure.com/",
+                mtls_proof_of_possession=True,
+                with_attestation_support=False,
+            )
+        mock_attest.assert_not_called()
+        self.assertEqual(result["access_token"], access_token)
+
 
 # ---------------------------------------------------------------------------
 # Attestation module