unity_integration.py

import os, random, string
from authlib.integrations.flask_client import OAuth
from flask import url_for, redirect, current_app as app
from flask_login import login_user
from flask_appbuilder import expose, BaseView as AppBuilderBaseView
from airflow.plugins_manager import AirflowPlugin
import logging
import os

log = logging.getLogger(__name__)
log.setLevel(os.getenv("AIRFLOW__LOGGING__FAB_LOGGING_LEVEL", "INFO"))

FAB_ADMIN_ROLE = "Admin"
FAB_VIEWER_ROLE = "Viewer"
FAB_PUBLIC_ROLE = "Public"  # The "Public" role is given no permissions

oauth = OAuth(app)
oauth.register(
    name='unity',
    client_id=os.getenv("OAUTH_CLIENT_ID"),
    server_metadata_url=os.getenv("OAUTH_METADATA_URL"),
    client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
    client_kwargs={
        'scope' : 'openid email profile eflows'
    }

)

class UnityIntegrationLoginView(AppBuilderBaseView):
    #@expose("/unity_login")
    @app.route('/unity_login')
    def unity_login():
        redirect_uri = url_for('unity_authorize', _external=True)
        return oauth.unity.authorize_redirect(redirect_uri)
class UnityIntegrationAuthView(AppBuilderBaseView):
    #@expose("/unity_authorize")
    @app.route('/unity_authorize')
    async def unity_authorize():
        token = await oauth.unity.authorize_access_token()
        user = await oauth.unity.userinfo(token=token)
        # get relevant data from token
        email = user['email']
        persistent_identifier = user["sub"]
        first_name = user["given_name"]
        last_name = user["family_name"]
        admin_access = user.get('eflows:dlsAccess', False)
        role = FAB_VIEWER_ROLE
        if admin_access:
            role = FAB_ADMIN_ROLE

        # check airflow user backend
        # check if user already exists, if not create it (with long random password)
        sec_manager = app.appbuilder.sm
        fab_user = sec_manager.find_user(username=persistent_identifier)
        if fab_user is None: # TODO check if None is the rioght thing to compare to
            characters = string.ascii_letters + string.digits + string.punctuation
            fab_user = sec_manager.add_user(
                username=persistent_identifier,
                first_name=first_name,
                last_name=last_name,
                email=email,
                role=role,
                password=''.join(random.choice(characters) for i in range(20))
            )
        # login as that user
        login_user(fab_user, remember=False)
        return redirect('/')

v_unity_login_view = UnityIntegrationLoginView()
v_unity_login_package = {
    "name": "Unity Login View",
    "category": "Unity Integration",
    "view": v_unity_login_view,
}

v_unity_auth_view = UnityIntegrationAuthView()
v_unity_auth_package = {
    "name": "Unity Auth View",
    "category": "Unity Integration",
    "view": v_unity_auth_view,
}

class UnityIntegrationPlugin(AirflowPlugin):
    name = "unity_integration"
    appbuilder_views = [v_unity_auth_package, v_unity_login_package]