Flask dance example for login with Azure AD

I am trying to implement SSO for one of my applications using flask-login and flask-dance. As a starting point I am using sample code given on Flask Dance website - https://flask-dance.readthedocs.io/en/v1.2.0/quickstarts/sqla-multiuser.html

Only change I did was - I replaced GitHub with my Azure AD credentials

Please find the code below:

import sys
from flask import Flask, redirect, url_for, flash, render_template
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.orm.exc import NoResultFound
from flask_dance.contrib.github import make_github_blueprint, github
from flask_dance.contrib.azure import make_azure_blueprint, azure
from flask_dance.consumer.storage.sqla import OAuthConsumerMixin, SQLAlchemyStorage
from flask_dance.consumer import oauth_authorized, oauth_error
from flask_login import (
    LoginManager, UserMixin, current_user,
    login_required, login_user, logout_user
)

# setup Flask application
app = Flask(__name__)
app.secret_key = "XXXXXXXXXXXXXX"
blueprint = make_azure_blueprint(
    client_id="XXXXXXXXXXXXXXXXXXXXX",
    client_secret="XXXXXXXXXXXXXXXXXXXXXXXX",
    tenant="XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
   
)
app.register_blueprint(blueprint, url_prefix="/login")

# setup database models
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///multi.db"
db = SQLAlchemy()

class User(db.Model, UserMixin):
    id = db.Column(db.Integer, primary_key=True)
    # Your User model can include whatever columns you want: Flask-Dance doesn't care.
    # Here are a few columns you might find useful, but feel free to modify them
    # as your application needs!
    username = db.Column(db.String(1028), unique=True)
    email = db.Column(db.String(1028), unique=True)
    name = db.Column(db.String(1028))

class OAuth(OAuthConsumerMixin, db.Model):
    provider_user_id = db.Column(db.String(1028), unique=True)
    user_id = db.Column(db.Integer, db.ForeignKey(User.id))
    user = db.relationship(User)

# setup login manager
login_manager = LoginManager()
login_manager.login_view = 'azure.login'

@login_manager.user_loader
def load_user(user_id):
    #print(User.query.get(int(user_id)))
    
    return User.query.get(int(user_id))

# setup SQLAlchemy backend
blueprint.storage = SQLAlchemyStorage(OAuth, db.session, user=current_user,user_required=False)

# create/login local user on successful OAuth login
@oauth_authorized.connect_via(blueprint)
def azure_logged_in(blueprint, token):
    if not token:
        #print(token)
        flash("Failed to log in with azure.", category="error")
        return False

    resp = blueprint.session.get("/user")
    if not resp.ok:
        #print(resp)
        msg = "Failed to fetch user info from Azure."
        flash(msg, category="error")
        return False

    azure_info = resp.json()
    azure_user_id = str(azure_info["id"])
    #print(azure_user_id)
    # Find this OAuth token in the database, or create it
    query = OAuth.query.filter_by(
        provider=blueprint.name,
        provider_user_id=azure_user_id,
    )
    try:
        oauth = query.one()
    except NoResultFound:
        oauth = OAuth(
            provider=blueprint.name,
            provider_user_id=azure_user_id,
            token=token,
        )

    if oauth.user:
        login_user(oauth.user)
        flash("Successfully signed in with Azure.")

    else:
        # Create a new local user account for this user
        user = User(
            # Remember that `email` can be None, if the user declines
            # to publish their email address on GitHub!
            email=azure_info["email"],
            name=azure_info["name"],
        )
        # Associate the new local user account with the OAuth token
        oauth.user = user
        # Save and commit our database models
        db.session.add_all([user, oauth])
        db.session.commit()
        # Log in the new local user account
        login_user(user)
        flash("Successfully signed in with Azure.")

    # Disable Flask-Dance's default behavior for saving the OAuth token
    return False

# notify on OAuth provider error
@oauth_error.connect_via(blueprint)
def azure_error(blueprint, error, error_description=None, error_uri=None):
    msg = (
        "OAuth error from {name}! "
        "error={error} description={description} uri={uri}"
    ).format(
        name=blueprint.name,
        error=error,
        description=error_description,
        uri=error_uri,
    )
    flash(msg, category="error")

@app.route("/logout")
@login_required
def logout():
    logout_user()
    flash("You have logged out")
    return redirect(url_for("index"))

@app.route("/")
def index():
    return render_template("home.html")

# hook up extensions to app
db.init_app(app)
login_manager.init_app(app)

if __name__ == "__main__":
    if "--setup" in sys.argv:
        with app.app_context():
            db.create_all()
            db.session.commit()
            print("Database tables created")
    else:
        app.run(debug=True,port=5011)

I have also done appropriate changes in HTML file for 'azure.login'. So after running it as python multi.py --setup database tables are getting created and after I run python multi.py Oauth dance is actually starting but in the end I am getting error like below:

HTTP Response:

127.0.0.1 - - [28/Oct/2020 10:17:44] "?[32mGET /login/azure/authorized?code=0.<Token>HTTP/1.1?[0m" 302 -
127.0.0.1 - - [28/Oct/2020 10:17:44] "?[37mGET / HTTP/1.1?[0m" 200 -

Am I missing something? Is it a good idea to use Flask Dance and Flask Login to have SSO with Azure AD? Or I should go with MSAL only along with Flask Session? Kindly give your valuable inputs..

Since you use Azure AD as the Flask dance provider, we need to use Microsoft Graph to get user's information. The URL should be https://graph.microsoft.com/v1.0/me. So please update the code resp = blueprint.session.get("/user") to resp = blueprint.session.get("/v1.0/me") in method azure_logged_in. Besides, please note that the azure ad user's information has different names. We also need to update the code about creating users.

For example

@oauth_authorized.connect_via(blueprint)
def azure_logged_in(blueprint, token):
    if not token:
        # print(token)
        flash("Failed to log in with azure.", category="error")
        return False

    resp = blueprint.session.get("/v1.0/me")
    # azure.get
    if not resp.ok:
        # print(resp)
        msg = "Failed to fetch user info from Azure."
        flash(msg, category="error")
        return False

    azure_info = resp.json()
    azure_user_id = str(azure_info["id"])
    # print(azure_user_id)
    # Find this OAuth token in the database, or create it
    query = OAuth.query.filter_by(
        provider=blueprint.name,
        provider_user_id=azure_user_id,
    )
    try:
        oauth = query.one()
    except NoResultFound:
        oauth = OAuth(
            provider=blueprint.name,
            provider_user_id=azure_user_id,
            token=token,
        )

    if oauth.user:
        login_user(oauth.user)
        flash("Successfully signed in with Azure.")

    else:
        # Create a new local user account for this user
        user = User(
            # create user with user information from Microsoft Graph
            email=azure_info["mail"],
            username=azure_info["displayName"],
            name=azure_info["userPrincipalName"]
        )
        # Associate the new local user account with the OAuth token
        oauth.user = user
        # Save and commit our database models
        db.session.add_all([user, oauth])
        db.session.commit()
        # Log in the new local user account
        login_user(user)
        flash("Successfully signed in with Azure.")

    # Disable Flask-Dance's default behavior for saving the OAuth token
    return False

For more details, please refer to here and here