JWT and Azure AD Authentication: How to Marry Them?

Hi everyone!

Today, I would like to highlight a really interesting topic: how to implement an additional authentication layer over a service that does not offer it out of the box. Namely, how to add Azure AD Authentication on your side over a service-side JWT token authentication.

It’s a good exercise in case your third-party service does not support SSO authentication (or you just don’t want to pay much for it) - but you would be happy with a simple solution. All you need is Python, Docker, and, surely, an understanding of what you are going to do. Don’t use this solution in production as it is :)

Let’s start!

Prerequisites

Before coding and deployment, let’s make sure we have everything we may need.

Third-party service to authenticate into

I assume that a destination service is able to do the following:

Generate a JWT secret to create JWT tokens;
Validate incoming requests with JWT tokens - and grant access in case the JWT token is valid;
If a JWT token is missing/invalid, redirect a user via a pre-defined fallback URL.

In my case, the service required a JWT token as a part of the URL to the content: https://dest-url.com/content/?jwt_token=<token>. So I simply had to build a redirect URL with a JWT token after all operations on my side. Also, JWT tokens can be passed in different ways, but anyway, it’s not rocket science to figure out how to include them in your resulting request.

Azure

To provide authentication, create a fresh application in your Azure organization. To control the authentication flow on your side, you will need the following details:

Application ID (in other words, Client ID);
Application secret;
Directory (tenant) ID;
List of application scopes (as a rule, Read permission will be enough);
Redirect URI.

The first four parameters will be generated when you create your application. As for the redirect URI, you have to provide it in advance: it will be the address of your Python service generating the authorization key for Azure.

Python service

We are going to connect JWT and AD authentication via a dockerized Python application deployed on our side. It will be based on Flask and the Microsoft Authentication Library for Python (MSAL). Below, you will find a list of vital Python libraries you will need:

flask: overall application architecture and endpoint design;
msal: authentication workflow;
jwt: generating JWT tokens based on third-party service JWT secrets;
fernet: additional encryption of a user-side session cookie.

Moreover, it would be really useful to get acquainted with the official Microsoft documentation.

Here, you will find a complete and more complicated example of a Flask-based application.

Authentication Workflow

The overall authentication workflow looks as follows:

The user goes to our https://python-authentication-service.com/. If they try to access the destination service directly, it will also redirect them for authentication via this fallback URL.
The authentication service initiates the AD authentication flow and invites the user to enter their credentials in the standard Microsoft authentication window. After successful login, the service acquires the authorization code.
The authentication flow creates a corresponding JSON object containing the authorization code acquired. The authentication service encrypts this object and saves it on the user side as a session cookie.
The authentication service decrypts the flow in the cookie and queries the AD application with the authorization code to generate an AD authentication token.
When the token is generated, the authentication service acquires the final authentication result.
If authentication is successful, the authentication service generates the JWT token based on a destination service JWT secret and redirects the user to the destination content with this token included.

If any of the tokens expire, the user will be redirected to either log in again or get a new destination service JWT token via the authentication service.

It is not straightforward why we need session cookies and why we encrypt them in steps 3 and 4. It is required only in case you want to deploy your service in a replicated manner: due to MSAL library constraints, the authentication service initiates the authentication flow and generates the AD tokens on two different endpoints. If you replicate the authentication service, your replicas will not be aware of any authentication flow initiated on another replica - so you need to save the flow object somewhere.

The easiest way to achieve this is to use the Flask session library and save all the stuff on the user side. In this case, all replicas of the authentication service will retain access to this shared cookie. Flow encryption here is used only to provide an additional security layer: since it’s a user-side cookie, we can rely on lightweight encryption algorithms.

Alternatively, if you deploy the service in a singleton pattern, you would not need session cookies. Moreover, you can play hard and create shared session storage on the server side, e.g., with the use of Redis or alternative solutions.

Implementation

I guess we’ve covered all theoretical steps and are ready for fun! Now, let’s implement the service.

Code

First of all, let’s define the overall project structure - it’s simple:

/src
├── application.py
└── config.py
Dockerfile
requirements.txt

Basically, project requirements should include the following libraries:

Flask>=2.3.2,<3
msal>=1.22.0,<2
requests>=2.31.0,<3
fernet>=1.0.1

In the configuration file, specify all required parameters - it’s still simple:

MSAL_CLIENT_ID = #Your AD application (client) ID
MSAL_CLIENT_SECRET = #Your AD application secret
MSAL_AUTHORITY = #https://login.microsoftonline.com/your-tenant-id
AD_SCOPES = #List of scopes defined when creating the AD application (for testing purposes, can be empty: '[]')

AUTH_SERVICE_URL = #URL of this authentication service (e.g. https://python-authentication-service.com/ as above)

JWT_SECRET = #JWT secret generated by the destination service
JWT_EXPIRATION_SECONDS = #Desired TTL for the destination service JWT token

FLASK_SECRET = #Built-in Flask cookie secret key to encrypt the session cookie with
FERNET_SECRET = #Additional Fernet secret key to encrypt the flow object

TARGET_URL = #URL of the destination service to redirect into with the resulting JWT token (e.g. https://dest-url.com/content/ as above)

You can store the configuration parameters as you wish. But since they are pretty sensitive, consider a secret manager to protect them.

Now, all preliminary stuff is done, and we are good to go ahead with the main code. Let’s import all required libraries and declare the Flask application, MSAL, and Fernet objects in use:

from flask import Flask, session, request, redirect
import msal
import jwt
from fernet import Fernet
import json
import time
import config

app = Flask(__name__)

msal_app = msal.ConfidentialClientApplication(
    client_id=config.MSAL_CLIENT_ID,
    authority=config.MSAL_AUTHORITY,
    client_credential=config.MSAL_CLIENT_SECRET,
)

fe = Fernet(bytes(config.FERNET_SECRET, encoding='utf8'))

Since we are going to use the Flask session library to store session cookies, we also have to define the cookie secret we added to the config before:

app.config.update(SECRET_KEY=config.FLASK_SECRET, ENV='development')

As I’ve already mentioned before, the application should support two endpoints:

/ - root endpoint to start the authentication flow and create the corresponding object with the authorization key.
/getadtoken - auxiliary authorization endpoint to validate authentication flow and create an AD authentication token. When done, this endpoint proceeds to the core stuff: generates the destination service JWT token, appends it to the destination URL, and redirects the user.

That’s all we need here. First, let’s implement the root endpoint:

@app.route('/')
def root():
    flow = msal_app.initiate_auth_code_flow(config.AD_SCOPES, redirect_uri=config.AUTH_SERVICE_URL + '/getadtoken')
    
    session['flow'] = fe.encrypt(json.dumps(flow).encode())

    return redirect(flow['auth_uri'], code=302)

Here you can see that we create the session object and add the encrypted flow inside. The session will be stored on the user side and shared between the application replicas, as I’ve mentioned before.

The flow object also contains the authentication URI on the Azure side to redirect the user for authentication.

Proceeding to the /getadtoken endpoint:

@app.route('/getadtoken')
def getadtoken():
    if 'flow' not in session:
        return 'Unauthorised', 403

    flow = json.loads(fe.decrypt(session['flow']).decode())
    
    saml_result = msal_app.acquire_token_by_auth_code_flow(flow, request.args)

    if not saml_result or 'error' in saml_result:
        return 'Unauthorised', 403

    jwt_token = jwt.encode({'data': 'empty', 'exp': int(time.time()) + config.JWT_EXPIRATION_SECONDS}, config.JWT_SECRET, algorithm='RS256') #NB: your encoding algorithm may differ.

    return redirect(config.TARGET_URL + '?jwt_token=' + jwt_token, code=302)

Here, first of all, we have to check if the flow object is inside the Flask session to prevent direct access to the /getadtoken endpoint. In other words, if the user does not have any session cookie stored, they are unauthenticated and cannot proceed.

Then, we decrypt the flow object and pass it to generate the AD authentication token. The authorization code contained in the flow is validated on the Azure side, so if there are any problems in the resulting object, the application returns 403.

If the authentication is successful, we finally proceed to the core stuff: generate the JWT token based on the destination service secret and let the user go ahead. Depending on your destination service, the JWT token may require any data inside, a particular expiration time to define, or a different encoding algorithm.

That’s all, folks! You are breathtaking!

Containerization

Since our application is Flask-based, you can use official Python base images to create yours:

FROM python:3.11.3-alpine

RUN apk update

WORKDIR /app

COPY requirements.txt ./

RUN pip install requirements.txt

COPY ./src .

CMD ["python", "./application.py"]

As I mentioned before, you can deploy the resulting service in any manner you wish.

What’s next?

As you can see, I’ve described the simplest possible approach to implement the Azure AD authentication on your side. It highlights the core principles and can be paired with many third-party services supporting JWT authentication. There are many ways to improve; let me show some to you:

Server-side session storage for better-replicated deployment;
Caching and rotation for both tokens;
User-side logout flow and endpoint.

Thank you for reading! Hope you’ve enjoyed this article.

Thank you in advance for your comments.