pythonoauthairflowamazon-cognitoflask-appbuilder

AWS Cognito OAuth configuration for Flask Appbuilder


I am setting up RBAC with Airflow, and testing locally to start. I have provisioned an AWS Cognito User Group via the console. Additionally, I have a webserver_config.py file I have mounted to my Airflow docker container to set up OAuth with RBAC.

Relevant section in my webserver_config.py file:

COGNITO_URL = os.getenv('COGNITO_URL')
CONSUMER_KEY = os.getenv('COGNITO_CLIENT_KEY')
SECRET_KEY = os.getenv('COGNITO_CLIENT_SECRET')

# When using OAuth Auth, uncomment to setup provider(s) info
# Google OAuth example:
OAUTH_PROVIDERS = [{
  'name':'AWS Cognito',
    'whitelist': ['@company.com'],  # optional
    'token_key':'access_token',
    'icon':'fa-amazon',
        'remote_app': {
            'base_url': os.path.join(COGNITO_URL, 'oauth2/idpresponse'),
            # 'base_url': COGNITO_URL,
            'request_token_params':{
                'scope': 'email profile'
            },
            'access_token_url': os.path.join(COGNITO_URL, 'oauth2/token'),
            'authorize_url': os.path.join(COGNITO_URL, 'oauth2/authorize'),
            'request_token_url': None,
            'consumer_key': CONSUMER_KEY,
            'consumer_secret': SECRET_KEY,
        }
}]

Variables are as follows:

COGNITO_URL: The domain name I have created in the "App Integration" section of my user pool

COGNITO_CLIENT_KEY: The app client id for my app in the "App Clients" section of my user pool

COGNITO_CLIENT_SECRET: The app client secret for my app in the "App Clients" section of my user pool

In the Cognito UI, I have the following settings for my App Client: enter image description here

Basically, I have set the endpoints as they should be on my local machine when testing. I have fiddled with both the http://localhost:8083/oauth2/idpresponse and http://localhost:8083/admin (normal home page for Airflow) routes and received the same error.

I think that the issue is that the URI the client is trying to request and the URI specified do not match. I tried following the advice at https://stackoverflow.com/a/53602884/13717098, but when I extracted that URI and saved it in the Cognito console, I continue to get the same error. I am looking for help identifying the URI needed. The request I've identified per the linked post is: /oauth2/authorize?response_type=code&client_id=269vguq386076suj80vpq4ctmj&redirect_uri=http%3A%2F%2Flocalhost%3A8083%2Foauth-authorized%2FAWS%2520Cognito&scope=email+profile&state=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJuZXh0IjpbImh0dHA6Ly9sb2NhbGhvc3Q6ODA4My9ob21lIl19.CcuxpZyuVIqW0GtnNL219Xkg1IftE0tzFiVilR6b4us I would appreciate any help with identifying the URI and/or its associated patterns.

Edited for spacing.


Solution

  • Flask builder library uses the name of the config object as value in redirect_uri.

    Set callback value to: http://localhost:8083/oauth-authorized/AWS%20Cognito instead of http://localhost:8080/oauth2/idresponse in AWS Cognito client. This should solve the redirection issue.

    The real problem will start for userinfo endpoint as AWS cognito uses OpenID auth pattern.

    aws-cognito-client

    EDIT

    AWS Cognito has oauth2/userinfo endpoint for receiving user information. To retrieve the userinfo, you're supposed to send openid scope along with your request. Following is my webserver_config.py.

    from airflow.www_rbac.security import AirflowSecurityManager
    from flask_appbuilder.security.manager import AUTH_OAUTH
    import os
    import json
    
    class CognitoSecurity(AirflowSecurityManager):
        def oauth_user_info(self, provider, response=None):
            if provider == "aws_cognito":
                me = self.appbuilder.sm.oauth_remotes[provider].get("userInfo")
                data = json.loads(me.raw_data)
                print("User info from aws_cognito: {0}".format(data))
                return {"username": data.get("username"), "email": data.get("email")}
            else:
                return {}
    
    AUTH_TYPE = AUTH_OAUTH
    
    AUTH_USER_REGISTRATION = True
    
    AUTH_USER_REGISTRATION_ROLE = "Admin"
    
    COGNITO_URL = ""
    CONSUMER_KEY = ""
    SECRET_KEY = ""
    
    OAUTH_PROVIDERS = [{
        'name':'aws_cognito',
        'whitelist': ['@positsource.com'],  # optional
        'token_key':'access_token',
        'url': COGNITO_URL,
        'icon': 'fa-amazon',
        'remote_app': {
            'base_url': os.path.join(COGNITO_URL, 'oauth2/idpresponse'),
            'request_token_params': {
                'scope': 'email profile openid'
            },
            'access_token_url': os.path.join(COGNITO_URL, 'oauth2/token'),
            'authorize_url': os.path.join(COGNITO_URL, 'oauth2/authorize'),
            'request_token_url': None,
            'consumer_key': CONSUMER_KEY,
            'consumer_secret': SECRET_KEY,
        }
    }]
    
    SECURITY_MANAGER_CLASS = CognitoSecurity
    

    This should get the airflow webserver working with AWS cognito. Roles and permissions management can be done by you.