Integration of AWS Cognito with Superset for OAuth Provider with RS256 Token Verification

[anuragsingh-prospera]

Hello everyone,

I've successfully integrated Superset with AWS Cognito as an OAuth provider. Users created in the Cognito user pool can log in to Superset. However, I'm facing an issue with generating the guest token. The access token generated by Cognito is encrypted using the RS256 algorithm, while Superset supports the HS256 algorithm by default. I found out that we can change the decoding algorithm by setting JWT_DECODE_ALGORITHMS = ['RS256']

Here are the changes I have made so far:

`from flask_appbuilder.security.manager import AUTH_OAUTH
from superset.security import SupersetSecurityManager, AuthOAuthView
from flask import redirect
from flask_appbuilder.security.views import AuthOAuthView
from jwt.algorithms import RSAAlgorithm
from cryptography.hazmat.primitives import serialization
import requests
import time
import json
import logging

logger = logging.getLogger(name)

class CustomSsoAuthOAuthView(AuthOAuthView):
@expose("/login/")
def login(self, provider="cognito"):
return super().login(provider=provider)

@expose("/logout/")
def logout(self, provider="cognito", register=None):
    provider_obj = self.appbuilder.sm.oauth_remotes[provider]
    url = ("logout?client_id={}&response_type={}&redirect_uri={}".format(
            provider_obj.client_id,
            provider_obj.server_metadata.get('response_type'),
            provider_obj.server_metadata.get('logout_redirect_uri')
            ))

    ret = super().logout()
    time.sleep(1)

    return redirect("{}{}".format(provider_obj.api_base_url, url))

class CustomSsoSecurityManager(SupersetSecurityManager):
global_public_key = None
authoauthview = CustomSsoAuthOAuthView

def oauth_user_info(self, provider, response=None):
    if provider == 'cognito':
        res = self.appbuilder.sm.oauth_remotes[provider].get('oauth2/userInfo')
        oauth_session = self.appbuilder.sm.oauth_remotes[provider].token
        logger.debug("OAuth session token: %s", oauth_session)
        access_token = oauth_session['access_token']
        id_token = oauth_session['id_token']
        jwks_url = 'https://cognito-idp.<region>.amazonaws.com/<regio>_<user pool id>/.well-known/jwks.json'
        jwks = requests.get(jwks_url).json()
        global_public_key = RSAAlgorithm.from_jwk(jwks['keys'][0])
        logger.debug("Global Public Key: %s", global_public_key)

        pem_public_key = global_public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        print("Hello world:", pem_public_key.decode('utf-8'))

        if res.raw.status != 200:
            logger.error('Failed to obtain user info: %s', res.data)
            return
        logger.debug("Full response content: %s", res._content)
        me = json.loads(res._content)
        logger.debug("User data: %s", me)
        username = me['username']
        email = me['email']
        first_name = me.get('given_name', '')
        last_name = me.get('family_name', '')
        return {
            'username': me['username'],
            'email': me['email'],
            'access_token': access_token,
        }

def auth_user_oauth(self, userinfo):
    logger.debug("User info: %s", userinfo)
    if not userinfo:
        logger.error("No user information received from OAuth provider.")
        return None

    username = userinfo.get('username')
    email = userinfo.get('email')
    access_token = userinfo.get('access_token')

    if not username or not email or not access_token:
        logger.error("Username, email, or token missing in user information.")
        return None
    
    user = self.find_user(username=username)
    if not user:
        logger.info(f"Creating new user: {username}")
        alpha_role = self.appbuilder.sm.find_role('Admin')
        logger.debug("Admin role: %s", alpha_role)
        try:
            user = self.add_user(
                username=username,
                first_name=username,
                last_name="",
                email=email,
                role=alpha_role,
            )
        except Exception as e:
            logger.error(f"Could not create user: {username}, error: {e}")
            return None

    self.update_user_auth_stat(user)
    return user

CUSTOM_SECURITY_MANAGER = CustomSsoSecurityManager

ENABLE_PROXY_FIX = True
AUTH_USER_REGISTRATION = True
AUTH_TYPE = AUTH_OAUTH
COGNITO_URL = ""
CLIENT_ID = ""
CLIENT_SECRET = ""
LOGOUT_REDIRECT_URI = "http://localhost:8088/login"

OAUTH_PROVIDERS = [{
'name': 'cognito',
'token_key': 'access_token',
'icon': 'fa-amazon',
'url': COGNITO_URL,
'remote_app': {
'client_id': CLIENT_ID,
'client_secret': CLIENT_SECRET,
'request_token_params': {
'scope': 'email openid profile'
},
'response_type': 'code',
'base_url': os.path.join(COGNITO_URL, 'oauth2/idpresponse'),
'access_token_url': os.path.join(COGNITO_URL, 'oauth2/token'),
'authorize_url': os.path.join(COGNITO_URL, 'oauth2/authorize'),
'access_token_method': 'POST',
'request_token_url': None,
'api_base_url': COGNITO_URL,
'logout_url': os.path.join(COGNITO_URL, 'logout'),
'logout_redirect_uri': LOGOUT_REDIRECT_URI
},
}]

JWT_DECODE_ALGORITHMS = ['RS256']
COGNITO_JWKS_URL = f' '
PUBLIC_KEY = requests.get(COGNITO_JWKS_URL).json()['keys'][0]

JWT_PUBLIC_KEY =PUBLIC_KEY
`
If anyone has suggestions on how to configure Superset to verify tokens generated by Cognito using the RS256 algorithm, I would greatly appreciate it.

Thank you!

[dosubot[bot]]

Hey @anuragsingh-prospera! I'm here to help you with your issue. Let's work together to solve this!

To configure Superset to verify tokens generated by AWS Cognito using the RS256 algorithm, you need to set up OAuth2 authentication in your superset_config.py file. Here is an example configuration:

1. Configure OAuth Provider:

from flask_appbuilder.security.manager import AUTH_OAUTH

# Set the authentication type to OAuth
AUTH_TYPE = AUTH_OAUTH

# Configure the OAuth provider for AWS Cognito
OAUTH_PROVIDERS = [
    {
        'name': 'cognito',
        'token_key': 'access_token',  # Name of the token in the response of access_token_url
        'icon': 'fa-address-card',    # Icon for the provider
        'remote_app': {
            'client_id': 'your_cognito_client_id',  # Client Id from AWS Cognito
            'client_secret': 'your_cognito_client_secret',  # Secret from AWS Cognito
            'client_kwargs': {
                'scope': 'openid profile email'  # Scopes for the Authorization
            },
            'access_token_method': 'POST',  # HTTP Method to call access_token_url
            'access_token_params': {
                'client_id': 'your_cognito_client_id'
            },
            'jwks_uri': 'https://cognito-idp.your-region.amazonaws.com/your_user_pool_id/.well-known/jwks.json',  # JWKS URI for token verification
            'api_base_url': 'https://cognito-idp.your-region.amazonaws.com/',
            'access_token_url': 'https://your_cognito_domain/oauth2/token',
            'authorize_url': 'https://your_cognito_domain/oauth2/authorize'
        }
    }
]

# Allow user self-registration
AUTH_USER_REGISTRATION = True

# Default user self-registration role
AUTH_USER_REGISTRATION_ROLE = "Public"

2. Create a Custom Security Manager:

import logging
from superset.security import SupersetSecurityManager

class CustomSsoSecurityManager(SupersetSecurityManager):
    def oauth_user_info(self, provider, response=None):
        logging.debug("Oauth2 provider: {0}.".format(provider))
        if provider == 'cognito':
            me = self.appbuilder.sm.oauth_remotes[provider].get('userinfo').data
            logging.debug("user_data: {0}".format(me))
            return {
                'name': me['name'],
                'email': me['email'],
                'id': me['sub'],
                'username': me['username'],
                'first_name': me['given_name'],
                'last_name': me['family_name']
            }

3. Add the Custom Security Manager to superset_config.py:

from custom_sso_security_manager import CustomSsoSecurityManager
CUSTOM_SECURITY_MANAGER = CustomSsoSecurityManager

4. Set JWT Decode Algorithms:

JWT_DECODE_ALGORITHMS = ['RS256']

This configuration will enable Superset to verify tokens generated by AWS Cognito using the RS256 algorithm ^[1].

To continue talking to Dosu, mention @dosu.