Apache Airflow Azure AAD SSO howto

apache_airflow_sso_howto.md

The following instructions for enabling Azure SSO for Apache Airflow nearly take you all the way - but fall short a couple of details around the configuration of airflow itself:

https://objectpartners.com/2021/12/24/enterprise-auth-for-airflow-azure-ad

All the "Azure" instructions there can be safely followed - the resulting webserver_config.py (which can be injected into a dockerised Airflow in /opt/airflow/webserver_config.py) can be built from the following:

from __future__ import annotations

import os

from airflow.www.fab_security.manager import AUTH_OAUTH
from airflow.www.security import AirflowSecurityManager
from airflow.utils.log.logging_mixin import LoggingMixin

basedir = os.path.abspath(os.path.dirname(__file__))

# Flask-WTF flag for CSRF
WTF_CSRF_ENABLED = True
WTF_CSRF_TIME_LIMIT = None

AUTH_TYPE = AUTH_OAUTH

OAUTH_PROVIDERS = [{
    'name':'Microsoft Azure AD',
    'token_key':'access_token',
    'icon':'fa-windows',
    'remote_app': {
        'api_base_url': "https://login.microsoftonline.com/{}".format(os.getenv("AAD_TENANT_ID")),
        'request_token_url': None,
        'request_token_params': {
            'scope': 'openid email profile'
        },
        'access_token_url': "https://login.microsoftonline.com/{}/oauth2/v2.0/token".format(os.getenv("AAD_TENANT_ID")),
        "access_token_params": {
            'scope': 'openid email profile'
        },
        'authorize_url': "https://login.microsoftonline.com/{}/oauth2/v2.0/authorize".format(os.getenv("AAD_TENANT_ID")),
        "authorize_params": {
            'scope': 'openid email profile'
        },
        'client_id': os.getenv("AAD_CLIENT_ID"),
        'client_secret': os.getenv("AAD_CLIENT_SECRET"),
        'jwks_uri': 'https://login.microsoftonline.com/common/discovery/v2.0/keys'
    }
}]

AUTH_USER_REGISTRATION_ROLE = "Public"
AUTH_USER_REGISTRATION = True
AUTH_ROLES_SYNC_AT_LOGIN = True
AUTH_ROLES_MAPPING = {
    "airflow_prod_admin": ["Admin"],
    "airflow_prod_user": ["Op"],
    "airflow_prod_viewer": ["Viewer"]
}

class AzureCustomSecurity(AirflowSecurityManager, LoggingMixin):
    def get_oauth_user_info(self, provider, response=None):
        me = self._azure_jwt_token_parse(response["id_token"])
        return {
            "name": me["name"],
            "email": me["email"],
            "first_name": me["given_name"],
            "last_name": me["family_name"],
            "id": me["oid"],
            "username": me["preferred_username"],
            "role_keys": me["roles"]
        }

# the first of these two appears to work with older Airflow versions, the latter newer.
FAB_SECURITY_MANAGER_CLASS = 'webserver_config.AzureCustomSecurity'
SECURITY_MANAGER_CLASS = AzureCustomSecurity

The above assumes environment variables are configured for the OAuth client secret, etc - and has been tested thoroughly and confirmed working.

Note the roles need to match what you configured in Azure (the example above is using airflow_prod_user etc, in deviation to the linked article above).

Hello @ctrongminh, I am using the configuration that you've provided in the code snippet. I am configuring airflow 2.10.3 in docker and I am not getting any payload/response through which I can extract the user details.

`from future import annotations
import os
from airflow.www.fab_security.manager import AUTH_OAUTH

from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride
from airflow.utils.log.logging_mixin import LoggingMixin

basedir = os.path.abspath(os.path.dirname(file))

WTF_CSRF_ENABLED = True
WTF_CSRF_TIME_LIMIT = None

AAD_TENANT_ID = "" # Tenant ID for your Azure Active Directory
AAD_CLIENT_ID = "" # Application (Client) ID
AAD_CLIENT_SECRET = "" # Application Client Secret

AUTH_TYPE = AUTH_OAUTH

OAUTH_PROVIDERS = [{
'name': f"{app-name}", # Replace with your app name
'token_key': 'access_token',
'icon': 'fa-sign-in',
'remote_app': {
'api_base_url': f"https://login.microsoftonline.com/{AAD_TENANT_ID}",
'request_token_url': None,
'request_token_params': {
'scope': 'openid email profile'
},
'access_token_url': f"https://login.microsoftonline.com/{AAD_TENANT_ID}/oauth2/v2.0/token",
'access_token_params': {
'scope': 'openid email profile'
},
'authorize_url': f"https://login.microsoftonline.com/{AAD_TENANT_ID}/oauth2/v2.0/authorize",
'authorize_params': {
'scope': 'openid email profile'
},
'client_id': f"", # Azure AD Client ID
'client_secret': f"", # Azure AD Client Secret
'jwks_uri': 'https://login.microsoftonline.com/{AAD_TENANT_ID}/discovery/v2.0/keys'
}
}]

AUTH_USER_REGISTRATION_ROLE = "ADSuper" # Custom Role created for new users
AUTH_USER_REGISTRATION = True
AUTH_ROLES_SYNC_AT_LOGIN = True

AUTH_ROLES_MAPPING = {
"Admin": ["Admin"],
"Op": ["Op"],
"Viewer": ["Viewer"],
"Public": ["Public"],
}

class AzureCustomSecurity(FabAirflowSecurityManagerOverride, LoggingMixin):
def get_oauth_user_info(self, provider, response=None):
self.log.debug(f"Parsing JWT token for provider : {provider}")

    try:
        # Get user info from OAuth provider
        me = super().get_oauth_user_info(provider, response)
        print(me)  # Debugging: print the response
        print(type(me))  # Debugging: print type of the response
        data = me.json()  # Extract JSON data from the response
        print("Sample payload is", data)  # Debugging: print the parsed payload

    except Exception as e:
        import traceback
        traceback.print_exc()  # Log the traceback for debugging
        self.log.debug(e)  # Log the error message

    # Log the parsed JWT token data for debugging
    self.log.debug(f"Parse JWT token : {me}")
    
    # Return user info in a dictionary format (mapping the attributes as needed)
    return {
        "username": data.get("sub", ""),  # Unique user identifier
        "first_name": data.get("given_name", ""),  # User's first name
        "last_name": data.get("family_name", ""),  # User's last name
        "email": data.get("email", ""),  # User's email address
        "role_keys": data.get("entitlement_group", []),  # User's roles/permissions
    }

Set the security manager class in Airflow config

FAB_SECURITY_MANAGER_CLASS = 'webserver_config.AzureCustomSecurity' # Old setting (still required in some cases)
SECURITY_MANAGER_CLASS = AzureCustomSecurity # New setting in Airflow 2.10.0
`

I am getting response as #Sample Payload is {}

There appears to be built in support for Azure AD / Entra ID auth.

If you look here, there is already an override defined for Azure:
https://github.com/apache/airflow/blob/be464b48d6b329a2510d312d7ff89f3c01d4e62b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py#L2206

So all that is required in webserver_config.py is the following:

Important to make sure that the provider name is "azure" as this is what determines how the token is handled in override.py.

from __future__ import annotations

import os

from airflow.www.fab_security.manager import AUTH_OAUTH

basedir = os.path.abspath(os.path.dirname(__file__))

AZURE_TENANT_ID = os.environ["AZURE_TENANT_ID"]
AZURE_CLIENT_ID = os.environ["AZURE_CLIENT_ID"]
AZURE_CLIENT_SECRET = os.environ["AZURE_CLIENT_SECRET"]

WTF_CSRF_ENABLED = True
WTF_CSRF_TIME_LIMIT = None

AUTH_TYPE = AUTH_OAUTH

# Auto create users
AUTH_USER_REGISTRATION = True

OAUTH_PROVIDERS = [
    {
        "name": "azure",
        "token_key": "id_token",
        "icon": "fa-windows",
        "remote_app": {
            "api_base_url": f"https://login.microsoftonline.com/{AZURE_TENANT_ID}/oauth2",
            "access_token_url": f"https://login.microsoftonline.com/{AZURE_TENANT_ID}/oauth2/v2.0/token",
            "authorize_url": f"https://login.microsoftonline.com/{AZURE_TENANT_ID}/oauth2/v2.0/authorize",
            "jwks_uri": f"https://login.microsoftonline.com/{AZURE_TENANT_ID}/discovery/v2.0/keys",
            "request_token_url": None,
            "client_id": AZURE_CLIENT_ID,
            "client_secret": AZURE_CLIENT_SECRET,
            "client_kwargs": {
                "scope": "User.read openid email profile",
                "resource": AZURE_CLIENT_ID,
            },
        },
    }
]

AUTH_ROLES_SYNC_AT_LOGIN = True
PERMANENT_SESSION_LIFETIME = 1800

AUTH_ROLES_MAPPING = {
    "AIRFLOW_VIEWER": ["Viewer"],
    "AIRFLOW_USER": ["User"],
    "AIRFLOW_ADMIN": ["Admin"],
}

Hi all,

Thanks for the insights!
I'm trying to implement this where the client secret is a certificate. The tenant id and client id are as expected. However, injecting the certificate as a client secret doesn't result in the correct response.

Is there an extra step needed to use a certificate as the client secret?