How to verify JWT produced by Azure Ad?

When I receive a JWK from Azure AD in Python, I would like to validate and decode it. I, however, keep getting the error "Signature verification failed".

I have the following setup:

1. Azure Setup
   In Azure I have created an app registration with the setting "Personal Microsoft accounts only".
2. Python Setup
   In Python I use the MSAL package for receiving tokens. And I use a public key from Azure to verify the token.

Using the credentials from the Azure Portal I set up a client for getting tokens.

import msal
ad_auth_client = msal.ConfidentialClientApplication(
    client_id = client_id,
    client_credential = client_secret,
    authority = "https://login.microsoftonline.com/consumers"
)
my_token = ad_auth_client.acquire_token_for_client(scopes=['https://graph.microsoft.com/.default'])

If I throw the token into a site like https://jwt.io/ everything looks good. Next I need public keys from Azure for verifying the token.

import requests
response = requests.get("https://login.microsoftonline.com/common/discovery/keys")
keys = response.json()['keys']

To match up the public key to the token, I use the 'kid' in the token header. I also get which algorithm was used for encryption.

import jwt
token_headers = jwt.get_unverified_header(my_token['access_token'])
token_alg = token_headers['alg']
token_kid = token_headers['kid']
public_key = None
for key in keys:
    if key['kid'] == token_kid:
        public_key = key

Now I have the correct public key from Azure to verify my token, but the problem is that it is a JWT key. Before I can use it for decoding, I need to convert it to a RSA PEM key.

from cryptography.hazmat.primitives import serialization
rsa_pem_key = jwt.algorithms.RSAAlgorithm.from_jwk(json.dumps(public_key))
rsa_pem_key_bytes = rsa_pem_key.public_bytes(
    encoding=serialization.Encoding.PEM, 
    format=serialization.PublicFormat.SubjectPublicKeyInfo
)
    

This is what the Azure Public Key looks like:

b'-----BEGIN PUBLIC KEY-----\nMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAyr3v1uETrFfT17zvOiy0\n1w8nO+1t67cmiZLZxq2ISDdte9dw+IxCR7lPV2wezczIRgcWmYgFnsk2j6m10H4t\nKzcqZM0JJ/NigY29pFimxlL7/qXMB1PorFJdlAKvp5SgjSTwLrXjkr1AqWwbpzG2\nyZUNN3GE8GvmTeo4yweQbNCd+yO/Zpozx0J34wHBEMuaw+ZfCUk7mdKKsg+EcE4Z\nv0Xgl9wP2MpKPx0V8gLazxe6UQ9ShzNuruSOncpLYJN/oQ4aKf5ptOp1rsfDY2IK\n9frtmRTKOdQ+MEmSdjGL/88IQcvCs7jqVz53XKoXRlXB8tMIGOcg+ICer6yxe2it\nIQIDAQAB\n-----END PUBLIC KEY-----\n'

The last thing I need to do is to verify the token using the public key.

decoded_token = jwt.decode(
    my_token['access_token'], 
    key=rsa_pem_key_bytes,
    verify=True,
    algorithms=[token_alg],
    audience=[client_id],
    issuer="https://login.microsoftonline.com/consumers"
)

The result I get is:

jwt.exceptions.InvalidSignatureError: Signature verification failed

I also tried to follow this popular guide: How to verify JWT id_token produced by MS Azure AD? Placing the x5c into the certificate pre- and postfixes just generated errors of invalid formatting.

Can you guys see any obvious errors? My main guess is that there is something wrong with the audience or the issuer, but I cannot pin down what it is, and Microsoft's documentation is horrible as always. Also, there is a secret key in the app registration in Azure, but it does not seem to work either.

So it turns out that my verification code was correct, but that I was trying to verify the wrong token. After creating slight modifications I now receive an id_token, which can be decoded and verified.