0

I am deploying Airflow 2.0 on Azure and using Azure AD for Authentication and Authorization. Created App registration and Custom roles. Authentication with AD works. need to read role claims in token to do the role mapping .. Below is my web Config ... For this I am overriding oauth_user_info method in AirflowSecurityManager to read the roles claim in Azure AD token and assign it to role_keys.. Problem is when I do this and assign 'SECURITY_MANAGER_CLASS = AzureCustomSecurity, although web UI POD gets up without any error but when I test it appears that custom class method oauth_user_info is not called at all ... if I remove the custom class then I can see authentication works and I can see the token in logs.... so problem is that this custom class in not invoked in the authentication & Authorization process. Any one can help here ?

I can see a similar kind of thing working for AWS congnito in a post here AWS Cognito OAuth configuration for Flask Appbuilder

This is my Web config

import os
import json
from airflow.configuration import conf
from flask_appbuilder.security.manager import AUTH_OAUTH
from airflow.www.security import AirflowSecurityManager
from customsecmanager import AzureDlabSecurity
SQLALCHEMY_DATABASE_URI = conf.get("core", "SQL_ALCHEMY_CONN")
basedir = os.path.abspath(os.path.dirname(__file__))
CSRF_ENABLED = True
AUTH_TYPE = AUTH_OAUTH
AUTH_USER_REGISTRATION_ROLE = "Public"
AUTH_USER_REGISTRATION = True

class AzureCustomSecurity(AirflowSecurityManager):
  def oauth_user_info(self, provider, response=None):
    log.debug("inside custom method")
    if provider == "azure":
        log.debug("Azure response received : {0}".format(resp))
        id_token = resp["id_token"]
        log.debug(str(id_token))
        me = self._azure_jwt_token_parse(id_token)
        log.debug("Parse JWT token : {0}".format(me))
        return {
            "name": me["name"],
            "email": me["upn"],
            "first_name": me["given_name"],
            "last_name": me["family_name"],
            "id": me["oid"],
            "username": me["oid"],
            "role_keys": me["roles"],       
          }    
    else:
        return {}


OAUTH_PROVIDERS = [{
    'name':'azure',
    'token_key':'access_token',
    'icon':'fa-windows',
    'remote_app': {
        'base_url':'https://graph.microsoft.com/v1.0/',
        'request_token_params' :{'scope': 'openid'},
        'access_token_url':'https://login.microsoftonline.com/<tenantid>/oauth2/token',
        'authorize_url':'https://login.microsoftonline.com/<tenantid>/oauth2/authorize',
        'request_token_url': None,
        'client_id':'XXXXXXXX',
        'client_secret':'******'
        }
}]

# a mapping from the values of `userinfo["role_keys"]` to a list of FAB roles
AUTH_ROLES_MAPPING = {
     "Viewer": ["Viewer"],
     "Admin": ["Admin"],
#}    
AUTH_ROLES_SYNC_AT_LOGIN = True
SECURITY_MANAGER_CLASS = AzureCustomSecurity
Manish Khare
  • 1
  • 1
  • 1
  • Did you ever find a way to integrate that? I kind of have a similar issue, I wrote an authentication method for Airflow behind Google IAP, now I cannot integrate it back. – Overbryd Apr 19 '21 at 19:41
  • I found this implementation from Astronomer, based on JWT tokens from an external Proxy, quite helpful: https://github.com/astronomer/astronomer-fab-securitymanager/blob/master/src/astronomer/flask_appbuilder/security.py – Overbryd Apr 19 '21 at 20:34

1 Answers1

0

I tried to reproduce your code, but instead of overriding oauth_user_info I've overridden get_oauth_user_info method (docs link) and it worked. May you try this?

class AzureCustomSecurity(AirflowSecurityManager):
  def get_oauth_user_info(self, provider, resp):
     """ you code here... """
vilozio
  • 111
  • 1
  • 8