Skip to content

Instantly share code, notes, and snippets.

@wallyhall

wallyhall/apache_airflow_sso_howto.md

Last active October 26, 2024 18:18
Show Gist options
  • Save wallyhall/915fedb4dfc766b61f442a32c95e1c29 to your computer and use it in GitHub Desktop.
Save wallyhall/915fedb4dfc766b61f442a32c95e1c29 to your computer and use it in GitHub Desktop.
Download ZIP
Apache Airflow Azure AAD SSO howto
Raw
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).

@nmaster
Copy link

nmaster commented Sep 13, 2024

Thanks a lot for this solution! Works fine for me in a PoC scenario using kubernetes 1.28 and the official airflow helm chart for version 2.9.3. I'm using this with Azure AD. The only flaw it has for me is that JWT2 doesn't come with first_name and last_name fields, however airflow seems to expect JWT1. My current airflow/flask/python skills surely are too limited, maybe somebody can help me out? ;)

sidenote: the described solution also supports multiple oauth providers. The only thing to make sure is that name of the oauth provider oapname (in this example 'azure') matches the redirect uri https://my.airflow.host/oauth-authorized/`oapname` configured on provider side.

@Saksham-lumiq
Copy link

@drivard
I have followed all the steps, but when I access my airflow, i still see the default airflow authentication method, although the URI changes but nothing else happens, its not redirecting or letting me authenticate using sso, any more steps you followed which missed out in the above conversations?
would really appreciate if some one could help me fix this.

@nitinmahawadiwar
Copy link

nitinmahawadiwar commented Oct 24, 2024 •
edited
Loading

Hello @drivard @nmaster @ctrongminh
I am using Airflow 2.9.3 with Helm 1.15
I have followed your procedure, but stuck with below error from the webserver pod.

Error authorizing OAuth access token: Invalid JSON Web Key Set.

This is what is reflecting in the UI : The request to sign in was denied

This is what I have configured in Azure Enterprise App:
Identifier (Entity ID) :: https://airflow.xyz.com/
Reply URL (Assertion Consumer Service URL) :: https://airflow.xyz.com/oauth-authorized/azure
Sign on URL :: https://airflow.xyz.com/login/
Relay State (Optional) :: https://airflow.xyz.com/home
Logout Url (Optional) :: https://airflow.xyz.com/logout

Below is my webserver_config.py

        from __future__ import annotations
        import os
        from airflow.www.fab_security.manager import AUTH_OAUTH
        # from airflow.www.security import AirflowSecurityManager
        from airflow.auth.managers.fab.security_manager.override import FabAirflowSecurityManagerOverride
        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
        AAD_TENANT_ID = <tenant id>
        AAD_CLIENT_ID = <APP Registration client id>
        AAD_CLIENT_SECRET = <App Registration client secret>

        AUTH_TYPE = AUTH_OAUTH

        OAUTH_PROVIDERS = [{
            'name':'azure',
            'token_key':'access_token',
            'icon':'fa-windows',
            'remote_app': {
                'api_base_url': f"https://login.microsoftonline.com/{AAD_TENANT_ID}",
                'client_kwargs': {
                    "scope": "User.read name preferred_username email profile upn",
                    "resource": f"{AAD_CLIENT_ID}",
                    # Optionally enforce signature JWT verification
                    "verify_signature": False
                },            
                '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"{AAD_CLIENT_ID}",
                'client_secret': f"{AAD_CLIENT_SECRET}",
                'jwks_uri': 'https://login.microsoftonline.com/common/discovery/v2.0/keys',
                'redirect_uri': 'https://airflow.xyz.com/oauth-authorized/azure'            
            }
        }]

        AUTH_USER_REGISTRATION_ROLE = "Public"
        AUTH_USER_REGISTRATION = True
        AUTH_ROLES_SYNC_AT_LOGIN = True
        # First you MUST create a role like"Admin with value Admin" in the App Registration "App Roles" section in the Azure Portal under Microsoft Entra ID.
        # Then groups MUST be linked from the Microsoft Entra ID "Enterprise Application" section in the Azure Portal under the "Users and Groups" section.
        # Each groups or users MUST be assigned a role e.g.: Admin, Op, Viewer in the "Users and Groups"
        AUTH_ROLES_MAPPING = {
            "airflow_nonprod_admin": ["Admin"],
            "airflow_nonprod_op": ["Op"],
            "airflow_nonprod_viewer": ["Viewer"],
        }

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

                try:   # the try and except are optional - strictly you only need the me= line.
                    me = super().get_oauth_user_info(provider, response)
                except Exception as e:
                    import traceback
                    traceback.print_exc()
                    self.log.debug(e)

                self.log.debug(f"Parse JWT token : {me}")
                return {
                    "name": me["userprincipalname"],
                    "email": me["mail"],
                    "first_name": me["givenname"],
                    "last_name": me["surname"],
                    "id": me["userprincipalname"],
                    "username": me["givenname"],
                    "role_keys": me["groups"]
                }

        # 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

would really appreciate for your help

@vdozal
Copy link

vdozal commented Oct 25, 2024 via email

@nitinmahawadiwar
Copy link

Thank you @vdozal . This was really helpful.

I am using almost same configuration except for the environment variables (but I am ignoring it for now.) .
Can you please help me to understand an issue with role_keys given below?

Let me explain the flow. un-wanted code is removed for clarity.

....
AUTH_USER_REGISTRATION_ROLE = "Public"
....
AUTH_ROLES_MAPPING = {
        "Entra ID admin Role": ["Admin"],
        "Entra ID viewer Role": ["Viewer"],
        "Entra ID Op Role": ["Op"],
 }
...
  class AzureCustomSecurity(FabAirflowSecurityManagerOverride,.....
............
 return {
            ................
                "role_keys": me.get("role_keys", ["airflow_public"])
            }

My observation is, inside the return block, even through "role_keys" is having Admin/Ops role values from Azure , its always selecting the role assigned for AUTH_USER_REGISTRATION_ROLE = "XXX" and the UI renders accordingly.

Log statements for me object is showing all correct values that are received from Azure.

Is there anything I am missing in this case ? My understanding is, whatever the role_keys inside the return block is set, the same role should be applied while rendering the UI.

@yriveiro
Copy link

@vdozal loading the AAD values from a secret with this values and it works:

webserver:
  webserverConfigConfigMapName: webserver-config-custom
  env:
    - name: AAD_TENANT_ID
      valueFrom:
        secretKeyRef:
          name: airflow-aad
          key: aad_tenant_id
    - name: AAD_CLIENT_ID
      valueFrom:
        secretKeyRef:
          name: airflow-aad
          key: aad_client_id
    - name: AAD_CLIENT_SECRET
      valueFrom:
        secretKeyRef:
          name: airflow-aad
          key: aad_client_secret

@yriveiro
Copy link

Does anyone know how I can get the REST API working as well?

My users will not have user and password, which doesn't allow me to use a basic auth provider. session backend probably works in the browser because it does some magic with the cookies, but I would like to call the REST API using curl

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment