Get JWT from Azure with Flask

fetch-jwt-from-azure.py

# /// script
# requires-python = ">=3.9"
# dependencies = [
#     "flask",
#     "requests",
# ]
# ///
# You can execute the entire script using the following command:
# uv run main.py

# Create a client_secrets.json that looks like this, based on your Azure credentials.
# Don't forget to include the callback URL as your allowed redirect return URL.
# {
#     "web": {
#         "client_id": "CLIENT_ID",
#         "client_secret": "CLIENT_SECRET",
#         "auth_uri": "https://login.microsoftonline.com/TENANT_ID/oauth2/authorize",
#         "token_uri": "https://login.microsoftonline.com/TENANT_ID/oauth2/token",
#         "userinfo_uri": "https://login.microsoftonline.com/TENANT_ID/oauth2/userinfo",
#         "issuer": "https://login.microsoftonline.com/TENANT_ID",
#         "redirect_uri": "http://localhost:5000/callback"
#     }
# }

import json
import os
from flask import Flask, redirect, request, session, url_for
import requests

app = Flask(__name__)
app.secret_key = os.urandom(24)

# Load OIDC configuration from client_secrets.json
with open("client_secrets.json") as f:
    oidc_config = json.load(f)["web"]
client_id = oidc_config["client_id"]
client_secret = oidc_config["client_secret"]
auth_uri = oidc_config["auth_uri"]
token_uri = oidc_config["token_uri"]
redirect_uri = oidc_config["redirect_uri"]


@app.route("/")
def index():
    msg = request.args.get("message")
    if msg:
        return msg
    return 'Initiate auth flow: <a href="/authorize">Go</a>'


@app.route("/authorize")
def login():
    return redirect(
        f"{auth_uri}?response_type=code&client_id={client_id}&redirect_uri={redirect_uri}&scope=openid profile email"
    )


@app.route("/callback")
def callback():
    code = request.args.get("code")
    if code:
        # Exchange authorization code for JWT token
        token_response = requests.post(
            token_uri,
            data={
                "grant_type": "authorization_code",
                "code": code,
                "redirect_uri": redirect_uri,
                "client_id": client_id,
                "client_secret": client_secret,
            },
        )
        token_json = token_response.json()

        if "id_token" in token_json:
            with open("token.json", "w") as f:
                json.dump(token_json, f)
            session["jwt_token"] = token_json["id_token"]
            return redirect(
                url_for(
                    "index",
                    message="Token saved successfully, check token.json",
                )
            )
        else:
            return f"Failed to retrieve JWT: {token_json}", 400
    return "Login failed", 400


if __name__ == "__main__":
    app.run(debug=True)