[Identity] Update async cert credential algorithm (Azure#39761)

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py

@@ -61,6 +61,7 @@ def __init__(
             self._custom_cache = True
         else:
             self._custom_cache = False
+        self._is_adfs = self._tenant_id.lower() == "adfs"
 
     def _get_cache(self, **kwargs: Any) -> TokenCache:
         cache = self._cae_cache if kwargs.get("enable_cae") else self._cache
@@ -239,7 +240,16 @@ def _get_jwt_assertion_request(self, scopes: Iterable[str], assertion: str, **kw
 
     def _get_client_certificate_assertion(self, certificate: AadClientCertificate, **kwargs: Any) -> str:
         now = int(time.time())
-        header = json.dumps({"typ": "JWT", "alg": "RS256", "x5t": certificate.thumbprint}).encode("utf-8")
+        headers = {"typ": "JWT"}
+        if self._is_adfs:
+            # Maintain backwards compatibility with older versions of ADFS.
+            headers["alg"] = "RS256"
+            headers["x5t"] = certificate.thumbprint
+        else:
+            headers["alg"] = "PS256"
+            headers["x5t#S256"] = certificate.sha256_thumbprint
+
+        jwt_header = json.dumps(headers).encode("utf-8")
         payload = json.dumps(
             {
                 "jti": str(uuid4()),
@@ -250,8 +260,8 @@ def _get_client_certificate_assertion(self, certificate: AadClientCertificate, *
                 "exp": now + (60 * 30),
             }
         ).encode("utf-8")
-        jws = base64.urlsafe_b64encode(header) + b"." + base64.urlsafe_b64encode(payload)
-        signature = certificate.sign(jws)
+        jws = base64.urlsafe_b64encode(jwt_header) + b"." + base64.urlsafe_b64encode(payload)
+        signature = certificate.sign_ps256(jws) if not self._is_adfs else certificate.sign_rs256(jws)
         jwt_bytes = jws + b"." + base64.urlsafe_b64encode(signature)
         return jwt_bytes.decode("utf-8")
 

sdk/identity/azure-identity/azure/identity/_internal/aadclient_certificate.py

@@ -26,7 +26,9 @@ def __init__(self, pem_bytes: bytes, password: Optional[bytes] = None) -> None:
 
         cert = x509.load_pem_x509_certificate(pem_bytes, default_backend())
         fingerprint = cert.fingerprint(hashes.SHA1())  # nosec
+        sha256_fingerprint = cert.fingerprint(hashes.SHA256())
         self._thumbprint = base64.urlsafe_b64encode(fingerprint).decode("utf-8")
+        self._sha256_thumbprint = base64.urlsafe_b64encode(sha256_fingerprint).decode("utf-8")
 
     @property
     def thumbprint(self) -> str:
@@ -36,11 +38,36 @@ def thumbprint(self) -> str:
         """
         return self._thumbprint
 
-    def sign(self, plaintext: bytes) -> bytes:
+    @property
+    def sha256_thumbprint(self) -> str:
+        """The certificate's SHA256 thumbprint as a base64url-encoded string.
+
+        :rtype: str
+        """
+        return self._sha256_thumbprint
+
+    def sign_rs256(self, plaintext: bytes) -> bytes:
         """Sign bytes using RS256.
 
         :param bytes plaintext: Bytes to sign.
         :return: The signature.
         :rtype: bytes
         """
         return self._private_key.sign(plaintext, padding.PKCS1v15(), hashes.SHA256())
+
+    def sign_ps256(self, plaintext: bytes) -> bytes:
+        """Sign bytes using PS256.
+
+        :param bytes plaintext: Bytes to sign.
+        :return: The signature.
+        :rtype: bytes
+        """
+        hash_alg = hashes.SHA256()
+
+        # Note: For PS265, the salt length should match the hash output size, so we use the hash algorithm's
+        # digest_size property to get the correct value.
+        return self._private_key.sign(
+            plaintext,
+            padding.PSS(mgf=padding.MGF1(hash_alg), salt_length=hash_alg.digest_size),
+            hash_alg,
+        )

sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py

@@ -21,13 +21,13 @@ class CertificateCredential(AsyncContextManager, GetTokenMixin):
 
     :param str tenant_id: ID of the service principal's tenant. Also called its 'directory' ID.
     :param str client_id: The service principal's client ID
-    :param str certificate_path: Path to a PEM-encoded certificate file including the private key. If not provided,
-          `certificate_data` is required.
+    :param str certificate_path: Optional path to a certificate file in PEM or PKCS12 format, including the private
+        key. If not provided, **certificate_data** is required.
 
     :keyword str authority: Authority of a Microsoft Entra endpoint, for example 'login.microsoftonline.com',
           the authority for Azure Public Cloud (which is the default). :class:`~azure.identity.AzureAuthorityHosts`
           defines authorities for other clouds.
-    :keyword bytes certificate_data: The bytes of a certificate in PEM format, including the private key
+    :keyword bytes certificate_data: The bytes of a certificate in PEM or PKCS12 format, including the private key.
     :keyword password: The certificate's password. If a unicode string, it will be encoded as UTF-8. If the certificate
           requires a different encoding, pass appropriately encoded bytes instead.
     :paramtype password: str or bytes

sdk/identity/azure-identity/tests/test_certificate_credential.py

@@ -297,6 +297,56 @@ def validate_jwt(request, client_id, cert_bytes, cert_password, expect_x5c=False
     cert.public_key().verify(signature, signed_part.encode("utf-8"), padding.PKCS1v15(), hashes.SHA256())
 
 
+def validate_jwt_ps256(request, client_id, cert_bytes, cert_password, expect_x5c=False):
+    """Validate the request meets Microsoft Entra ID's expectations for a client credential grant using a certificate, as documented
+    at https://learn.microsoft.com/entra/identity-platform/certificate-credentials
+    """
+
+    try:
+        cert = x509.load_pem_x509_certificate(cert_bytes, default_backend())
+    except ValueError:
+        if cert_password:
+            if isinstance(cert_password, str):
+                cert_password = cert_password.encode("utf-8")
+        cert_bytes = load_pkcs12_certificate(cert_bytes, cert_password).pem_bytes
+        cert = x509.load_pem_x509_certificate(cert_bytes, default_backend())
+
+    # jwt is of the form 'header.payload.signature'; 'signature' is 'header.payload' signed with cert's private key
+    jwt = request.body["client_assertion"]
+    if isinstance(jwt, bytes):
+        jwt = jwt.decode("utf-8")
+    header, payload, signature = (urlsafeb64_decode(s) for s in jwt.split("."))
+    signed_part = jwt[: jwt.rfind(".")]
+
+    claims = json.loads(payload.decode("utf-8"))
+    assert claims["aud"] == request.url
+    assert claims["iss"] == claims["sub"] == client_id
+
+    deserialized_header = json.loads(header.decode("utf-8"))
+    assert deserialized_header["alg"] == "PS256"
+    assert deserialized_header["typ"] == "JWT"
+    if expect_x5c:
+        # x5c should have all the certs in the file, in order, in PEM format minus headers and footers
+        pem_lines = cert_bytes.decode("utf-8").splitlines()
+        header = "-----BEGIN CERTIFICATE-----"
+        assert len(deserialized_header["x5c"]) == pem_lines.count(header)
+
+        # concatenate the PEM file's certs, removing headers and footers
+        chain_start = pem_lines.index(header)
+        pem_chain_content = "".join(line for line in pem_lines[chain_start:] if not line.startswith("-" * 5))
+        assert "".join(deserialized_header["x5c"]) == pem_chain_content, "JWT's x5c claim contains unexpected content"
+    else:
+        assert "x5c" not in deserialized_header
+    assert urlsafeb64_decode(deserialized_header["x5t#S256"]) == cert.fingerprint(hashes.SHA256())  # nosec
+
+    cert.public_key().verify(
+        signature,
+        signed_part.encode("utf-8"),
+        padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=hashes.SHA256.digest_size),
+        hashes.SHA256(),
+    )
+
+
 @pytest.mark.parametrize("cert_path,cert_password", ALL_CERTS)
 @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
 def test_token_cache_persistent(cert_path, cert_password, get_token_method):

sdk/identity/azure-identity/tests/test_certificate_credential_async.py

@@ -16,7 +16,7 @@
 
 from helpers import build_aad_response, mock_response, Request, GET_TOKEN_METHODS
 from helpers_async import async_validating_transport, AsyncMockTransport
-from test_certificate_credential import ALL_CERTS, EC_CERT_PATH, PEM_CERT_PATH, validate_jwt
+from test_certificate_credential import ALL_CERTS, EC_CERT_PATH, PEM_CERT_PATH, validate_jwt, validate_jwt_ps256
 
 
 def test_non_rsa_key():
@@ -174,19 +174,22 @@ def test_requires_certificate():
 @pytest.mark.asyncio
 @pytest.mark.parametrize("cert_path,cert_password", ALL_CERTS)
 @pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
-async def test_request_body(cert_path, cert_password, get_token_method):
+@pytest.mark.parametrize("tenant_id", ("adfs", "tenant"))
+async def test_request_body(cert_path, cert_password, get_token_method, tenant_id):
     access_token = "***"
     authority = "authority.com"
     client_id = "client-id"
     expected_scope = "scope"
-    tenant_id = "tenant"
 
     async def mock_send(request, **kwargs):
         assert request.body["grant_type"] == "client_credentials"
         assert request.body["scope"] == expected_scope
 
         with open(cert_path, "rb") as cert_file:
-            validate_jwt(request, client_id, cert_file.read(), cert_password)
+            if tenant_id == "adfs":
+                validate_jwt(request, client_id, cert_file.read(), cert_password)
+            else:
+                validate_jwt_ps256(request, client_id, cert_file.read(), cert_password)
 
         return mock_response(json_payload={"token_type": "Bearer", "expires_in": 42, "access_token": access_token})