Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
oauth2.py281 linesDownload Raw Back to authenticate
1##########################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8##########################################################################9 10"""A blueprint module implementing the Oauth2 authentication."""11 12import config13 14from authlib.integrations.flask_client import OAuth15from flask import current_app, url_for, session, request,\16    redirect, Flask, flash17from flask_babel import gettext18from flask_security import login_user, current_user19from flask_security.utils import logout_user20 21from pgadmin.authenticate.internal import BaseAuthentication22from pgadmin.model import User23from pgadmin.tools.user_management import create_user24from pgadmin.utils.constants import OAUTH2, MessageType25from pgadmin.utils import PgAdminModule, get_safe_post_login_redirect, \26    get_safe_post_logout_redirect27from pgadmin.utils.csrf import pgCSRFProtect28from pgadmin.model import db29 30OAUTH2_LOGOUT = 'oauth2.logout'31OAUTH2_AUTHORIZE = 'oauth2.authorize'32 33 34class Oauth2Module(PgAdminModule):35    def register(self, app, options):36        # Do not look for the sub_modules,37        # instead call blueprint.register(...) directly38        super().register(app, options)39 40    def get_exposed_url_endpoints(self):41        return [OAUTH2_AUTHORIZE,42                OAUTH2_LOGOUT]43 44 45def init_app(app):46    MODULE_NAME = 'oauth2'47 48    blueprint = Oauth2Module(MODULE_NAME, __name__, static_url_path='')49 50    @blueprint.route('/authorize', endpoint="authorize",51                     methods=['GET', 'POST'])52    @pgCSRFProtect.exempt53    def oauth_authorize():54        auth_obj = session['auth_obj']55        auth_obj.set_current_source(auth_obj.source.get_source_name())56        status, msg = auth_obj.login()57        if status:58            session['auth_source_manager'] = auth_obj.as_dict()59            if 'auth_obj' in session:60                session.pop('auth_obj')61            return redirect(get_safe_post_login_redirect())62        if 'auth_obj' in session:63            session.pop('auth_obj')64        logout_user()65        flash(msg, MessageType.ERROR)66        return redirect(get_safe_post_login_redirect())67 68    @blueprint.route('/logout', endpoint="logout",69                     methods=['GET', 'POST'])70    @pgCSRFProtect.exempt71    def oauth_logout():72        logout_url = None73        if 'oauth2_logout_url' in session:74            logout_url = session['oauth2_logout_url']75 76        if not current_user.is_authenticated:77            return redirect(get_safe_post_logout_redirect())78        for key in list(session.keys()):79            session.pop(key)80 81        logout_user()82        if logout_url:83            id_token = session['oauth2_token'].get('id_token')84            return redirect(logout_url.format(85                redirect_uri=request.url_root,86                id_token=id_token))87        return redirect(get_safe_post_logout_redirect())88 89    app.register_blueprint(blueprint)90    app.login_manager.logout_view = OAUTH2_LOGOUT91 92 93class OAuth2Authentication(BaseAuthentication):94    """OAuth Authentication Class"""95 96    LOGOUT_VIEW = OAUTH2_LOGOUT97 98    oauth_obj = OAuth(Flask(__name__))99    oauth2_clients = {}100    oauth2_config = {}101    email_keys = ['mail', 'email']102 103    def __init__(self):104        for oauth2_config in config.OAUTH2_CONFIG:105 106            OAuth2Authentication.oauth2_config[107                oauth2_config['OAUTH2_NAME']] = oauth2_config108 109            OAuth2Authentication.oauth2_clients[110                oauth2_config['OAUTH2_NAME']111            ] = OAuth2Authentication.oauth_obj.register(112                name=oauth2_config['OAUTH2_NAME'],113                client_id=oauth2_config['OAUTH2_CLIENT_ID'],114                client_secret=oauth2_config['OAUTH2_CLIENT_SECRET'],115                access_token_url=oauth2_config['OAUTH2_TOKEN_URL'],116                authorize_url=oauth2_config['OAUTH2_AUTHORIZATION_URL'],117                api_base_url=oauth2_config['OAUTH2_API_BASE_URL'],118                client_kwargs={'scope': oauth2_config.get(119                    'OAUTH2_SCOPE', 'email profile'),120                    'verify': oauth2_config.get(121                    'OAUTH2_SSL_CERT_VERIFICATION', True)},122                server_metadata_url=oauth2_config.get(123                    'OAUTH2_SERVER_METADATA_URL', None)124            )125 126    def get_source_name(self):127        return OAUTH2128 129    def get_friendly_name(self):130        return self.oauth2_config[self.oauth2_current_client]['OAUTH2_NAME']131 132    def validate(self, form):133        return True, None134 135    def login(self, form):136        profile = self.get_user_profile()137        email_key = \138            [value for value in self.email_keys if value in profile.keys()]139        email = profile[email_key[0]] if (len(email_key) > 0) else None140 141        username = email142        username_claim = None143        if 'OAUTH2_USERNAME_CLAIM' in self.oauth2_config[144                self.oauth2_current_client]:145            username_claim = self.oauth2_config[146                self.oauth2_current_client147            ]['OAUTH2_USERNAME_CLAIM']148        if username_claim is not None:149            if username_claim in profile:150                username = profile[username_claim]151            else:152                error_msg = "The claim '%s' is required to login into " \153                    "pgAdmin. Please update your OAuth2 profile." % (154                        username_claim)155                current_app.logger.exception(error_msg)156                return False, gettext(error_msg)157        else:158            if not email or email == '':159                error_msg = "An email id or OAUTH2_USERNAME_CLAIM is" \160                    " required to login into pgAdmin. Please update your" \161                    " OAuth2 profile for email id or set" \162                    " OAUTH2_USERNAME_CLAIM config parameter."163                current_app.logger.exception(error_msg)164                return False, gettext(error_msg)165 166        additinal_claims = None167        if 'OAUTH2_ADDITIONAL_CLAIMS' in self.oauth2_config[168                self.oauth2_current_client]:169 170            additinal_claims = self.oauth2_config[171                self.oauth2_current_client172            ]['OAUTH2_ADDITIONAL_CLAIMS']173 174        # checking oauth provider userinfo response175        valid_profile, reason = self.__is_any_claim_valid(profile,176                                                          additinal_claims)177        current_app.logger.debug(f"profile claims: {profile}")178        current_app.logger.debug(f"reason: {reason}")179 180        # checking oauth provider idtoken claims181        id_token_claims = session.get('oauth2_token', {}).get('userinfo',{})182        valid_idtoken, reason = self.__is_any_claim_valid(id_token_claims,183                                                          additinal_claims)184        current_app.logger.debug(f"idtoken claims: {id_token_claims}")185        current_app.logger.debug(f"reason: {reason}")186 187        if not valid_profile and not valid_idtoken:188            return_msg = "The user is not authorized to login" \189                " based on your identity profile." \190                " Please contact your administrator."191            audit_msg = f"The authenticated user {username} is not" \192                " authorized to access pgAdmin based on OAUTH2 config. " \193                f"Reason: additional claim required {additinal_claims}, " \194                f"profile claims {profile}, idtoken cliams {id_token_claims}."195            current_app.logger.warning(audit_msg)196            return False, return_msg197 198        user, msg = self.__auto_create_user(username, email)199        if user:200            user = db.session.query(User).filter_by(201                username=username, auth_source=OAUTH2).first()202            current_app.login_manager.logout_view = \203                OAuth2Authentication.LOGOUT_VIEW204            current_app.logger.info(205                "OAUTH2 user {0} logged in.".format(username))206            return login_user(user), None207        return False, msg208 209    def get_user_profile(self):210        session['oauth2_token'] = self.oauth2_clients[211            self.oauth2_current_client].authorize_access_token()212 213        session['pass_enc_key'] = session['oauth2_token']['access_token']214 215        if 'OAUTH2_LOGOUT_URL' in self.oauth2_config[216                self.oauth2_current_client]:217            session['oauth2_logout_url'] = self.oauth2_config[218                self.oauth2_current_client]['OAUTH2_LOGOUT_URL']219 220        resp = self.oauth2_clients[self.oauth2_current_client].get(221            self.oauth2_config[222                self.oauth2_current_client]['OAUTH2_USERINFO_ENDPOINT'],223            token=session['oauth2_token']224        )225        resp.raise_for_status()226        return resp.json()227 228    def authenticate(self, form):229        self.oauth2_current_client = request.form['oauth2_button']230        redirect_url = url_for(OAUTH2_AUTHORIZE, _external=True)231 232        if self.oauth2_current_client not in self.oauth2_clients:233            return False, gettext(234                "Please set the configuration parameters properly.")235        return False, self.oauth2_clients[236            self.oauth2_current_client].authorize_redirect(redirect_url)237 238    def __auto_create_user(self, username, email):239        if config.OAUTH2_AUTO_CREATE_USER:240            user = User.query.filter_by(username=username,241                                        auth_source=OAUTH2).first()242            if not user:243                create_msg = ("Creating user {0} with email {1} "244                              "from auth source OAUTH2.")245                current_app.logger.info(create_msg.format(username,246                                                          email))247                return create_user({248                    'username': username,249                    'email': email,250                    'role': 2,251                    'active': True,252                    'auth_source': OAUTH2253                })254 255        return True, {'username': username}256 257    def __is_any_claim_valid(self, identity, additional_claims):258        if additional_claims is None:259            reason = "Additional claim config is None, no check to do."260            return (True, reason)261        if not isinstance(additional_claims, dict):262            reason = "Additional claim check config is not a dict."263            return (False, reason)264        if additional_claims.keys() is None:265            reason = "Additional claim check config dict is empty."266            return (False, reason)267        for key in additional_claims.keys():268            claim = identity.get(key)269            if claim is None:270                continue271            if not isinstance(claim, list):272                claim = [claim]273            authorized_claims = additional_claims.get(key)274            if not isinstance(authorized_claims, list):275                authorized_claims = [authorized_claims]276            if any(item in authorized_claims for item in claim):277                reason = "Claim match found. Authorized access."278                return (True, reason)279        reason = f"No match was found."280        return (False, reason)281 
codekingpro/portable-devtools · Team Ai