keycloak.py

from typing import List, Set

import logging
from superset.security import SupersetSecurityManager

log = logging.getLogger(__name__)

class KeycloakSecurityManager(SupersetSecurityManager):
    def oauth_user_info(self, provider, resp):
        import pprint
        pprint.pprint(resp)

        log.debug("Oauth2 provider: {0}.".format(provider))
        provider_instance = self.appbuilder.sm.oauth_remotes[provider]

        client_id = provider_instance.client_id
        
        logging.debug("provider.client_id: {0}".format(client_id))
        ret = {
            'username': resp['userinfo']['preferred_username'],
            'email': resp['userinfo']['email'],
            'role_keys': resp['userinfo']['resource_access'][client_id]['roles'],
        }

        return ret

    def get_roles_from_keys(self, role_keys: List[str]) -> Set[SupersetSecurityManager.role_model]:
        _roles = set()
        for role_key in role_keys:
            fab_role = self.find_role(role_key)            
            if fab_role:
                _roles.add(fab_role)
            else:
                log.warning(f"Can't find role: {role_key}")
        return _roles