codekingpro/portable-devtools

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 Authentication."""11 12import config13import copy14import functools15from threading import Lock16 17from flask import current_app, flash, Response, request, url_for, \18    session, redirect, render_template19from flask_babel import gettext20from flask_security.views import _security, _ctx21from flask_security.utils import logout_user, config_value22 23from flask_login import current_user24from flask_socketio import disconnect, ConnectionRefusedError25 26from pgadmin.model import db, User27from pgadmin.utils.constants import KERBEROS, INTERNAL, OAUTH2, LDAP,\28    MessageType29import pgadmin.utils as pga_utils30from pgadmin.authenticate.registry import AuthSourceRegistry31 32MODULE_NAME = 'authenticate'33auth_obj = None34 35_URL_WITH_NEXT_PARAM = "{0}?next={1}"36 37 38class AuthLocker:39    """Implementing lock while authentication."""40    lock = Lock()41 42    def __enter__(self):43        self.lock.acquire()44        return self45 46    def __exit__(self, type, value, traceback):47        if self.lock.locked():48            self.lock.release()49 50 51def get_logout_url() -> str:52    """53    Returns the logout url based on the current authentication method.54 55    Returns:56        str: logout url57    """58    BROWSER_INDEX = 'browser.index'59    if config.SERVER_MODE and\60            session['auth_source_manager']['current_source'] == \61            KERBEROS:62        return _URL_WITH_NEXT_PARAM.format(url_for(63            'kerberos.logout'), url_for(BROWSER_INDEX))64    elif config.SERVER_MODE and\65            session['auth_source_manager']['current_source'] == \66            OAUTH2:67        return _URL_WITH_NEXT_PARAM.format(url_for(68            'oauth2.logout'), url_for(BROWSER_INDEX))69 70    return _URL_WITH_NEXT_PARAM.format(71        url_for('security.logout'), url_for(BROWSER_INDEX))72 73 74def socket_login_required(f):75    @functools.wraps(f)76    def wrapped(*args, **kwargs):77        if not current_user.is_authenticated:78            disconnect()79            raise ConnectionRefusedError("Unauthorised !")80        else:81            return f(*args, **kwargs)82    return wrapped83 84 85class AuthenticateModule(pga_utils.PgAdminModule):86    def get_exposed_url_endpoints(self):87        return ['authenticate.login']88 89 90blueprint = AuthenticateModule(MODULE_NAME, __name__, static_url_path='')91 92 93@blueprint.route('/login', endpoint='login', methods=['GET', 'POST'])94def login():95    """96    Entry point for all the authentication sources.97    The user input will be validated and authenticated.98    """99    with AuthLocker():100        return _login()101 102 103def _login():104    """105    Internal authentication process locked by a mutex.106    """107    form = _security.forms.get('login_form').cls(request.form)108    if OAUTH2 in config.AUTHENTICATION_SOURCES \109            and 'oauth2_button' in request.form:110        # Sending empty form as oauth2 does not require form attribute111        auth_obj = AuthSourceManager({}, copy.deepcopy(112            config.AUTHENTICATION_SOURCES))113        session['auth_obj'] = auth_obj114    else:115        auth_obj = AuthSourceManager(form, copy.deepcopy(116            config.AUTHENTICATION_SOURCES))117 118    session['auth_source_manager'] = None119 120    username = form.data['email']121    user = User.query.filter_by(username=username,122                                auth_source=INTERNAL).first()123 124    if user:125        if user.login_attempts >= config.MAX_LOGIN_ATTEMPTS > 0:126            user.locked = True127        else:128            user.locked = False129        db.session.commit()130 131        if user.login_attempts >= config.MAX_LOGIN_ATTEMPTS > 0:132            flash(gettext('Your account is locked. Please contact the '133                          'Administrator.'),134                  MessageType.WARNING)135            logout_user()136            return redirect(pga_utils.get_safe_post_logout_redirect())137 138    # Validate the user139    if not auth_obj.validate():140        for field in form.errors:141            flash_login_attempt_error = None142            if user and field in config.LOGIN_ATTEMPT_FIELDS:143                if config.MAX_LOGIN_ATTEMPTS > 0:144                    user.login_attempts += 1145                    left_attempts = \146                        config.MAX_LOGIN_ATTEMPTS - user.login_attempts147                    if left_attempts > 1:148                        flash_login_attempt_error = \149                            gettext('{0} more attempts remaining.'.150                                    format(left_attempts))151                    else:152                        flash_login_attempt_error = \153                            gettext('{0} more attempt remaining.'.154                                    format(left_attempts))155                db.session.commit()156            for error in form.errors[field]:157                if flash_login_attempt_error:158                    error = error + flash_login_attempt_error159                    flash_login_attempt_error = None160                flash(error, MessageType.WARNING)161 162        return redirect(pga_utils.get_safe_post_logout_redirect())163 164    # Authenticate the user165    status, msg = auth_obj.authenticate()166    if status:167        # Login the user168        status, msg = auth_obj.login()169        current_auth_obj = auth_obj.as_dict()170 171        if not status:172            if current_auth_obj['current_source'] == \173                    KERBEROS:174                return redirect('{0}?next={1}'.format(url_for(175                    'authenticate.kerberos_login'), url_for('browser.index')))176 177            flash(msg, MessageType.ERROR)178            return redirect(pga_utils.get_safe_post_logout_redirect())179 180        session['auth_source_manager'] = current_auth_obj181 182        if user:183            user.login_attempts = 0184        db.session.commit()185 186        if 'auth_obj' in session:187            session.pop('auth_obj')188        return redirect(pga_utils.get_safe_post_login_redirect())189 190    elif isinstance(msg, Response):191        return msg192    elif 'oauth2_button' in request.form and not isinstance(msg, str):193        return msg194    if 'auth_obj' in session:195        session.pop('auth_obj')196    flash(msg, MessageType.ERROR)197    form_class = _security.forms.get('login_form').cls198    form = form_class()199 200    return _security.render_template(201        config_value('LOGIN_USER_TEMPLATE'),202        login_user_form=form, **_ctx('login'))203 204 205class AuthSourceManager:206    """This class will manage all the authentication sources.207     """208 209    def __init__(self, form, sources):210        self.form = form211        self.auth_sources = sources212        self.source = None213        self.source_friendly_name = INTERNAL214        self.current_source = INTERNAL215        self.update_auth_sources()216 217    def as_dict(self):218        """219        Returns the dictionary object representing this object.220        """221 222        res = dict()223        res['source_friendly_name'] = self.source_friendly_name224        res['auth_sources'] = self.auth_sources225        res['current_source'] = self.current_source226 227        return res228 229    def update_auth_sources(self):230        for auth_src in [KERBEROS, OAUTH2]:231            if auth_src in self.auth_sources:232                if 'internal_button' in request.form:233                    self.auth_sources.remove(auth_src)234                else:235                    if INTERNAL in self.auth_sources:236                        self.auth_sources.remove(INTERNAL)237                    if LDAP in self.auth_sources:238                        self.auth_sources.remove(LDAP)239 240    def set_current_source(self, source):241        self.current_source = source242 243    @property244    def get_current_source(self):245        return self.current_source246 247    def set_source(self, source):248        self.source = source249 250    @property251    def get_source(self):252        return self.source253 254    def set_source_friendly_name(self, name):255        self.source_friendly_name = name256 257    @property258    def get_source_friendly_name(self):259        return self.source_friendly_name260 261    def validate(self):262        """Validate through all the sources."""263        err_msg = None264        for src in self.auth_sources:265            source = get_auth_sources(src)266            status, err_msg = source.validate(self.form)267            if status:268                return True269        if err_msg:270            flash(err_msg, MessageType.WARNING)271        return False272 273    def authenticate(self):274        """Authenticate through all the sources."""275        status = False276        msg = None277        for src in self.auth_sources:278            source = get_auth_sources(src)279            self.set_source(source)280            current_app.logger.debug(281                "Authentication initiated via source: %s" %282                source.get_source_name())283 284            status, msg = source.authenticate(self.form)285 286            if status:287                self.set_current_source(source.get_source_name())288                if msg is not None and 'username' in msg:289                    self.form._fields['email'].data = msg['username']290                return status, msg291            else:292                current_app.logger.debug(293                    "Authentication initiated via source: %s is failed." %294                    source.get_source_name())295 296        return status, msg297 298    def login(self):299        status, msg = self.source.login(self.form)300        if status:301            self.set_source_friendly_name(self.source.get_friendly_name())302            current_app.logger.debug(303                "Authentication and Login successfully done via source : %s" %304                self.source.get_source_name())305 306            # Set the login, logout view as per source  if available307            current_app.login_manager.login_view = getattr(308                self.source, 'LOGIN_VIEW', 'security.login')309            current_app.login_manager.logout_view = getattr(310                self.source, 'LOGOUT_VIEW', 'security.logout')311 312        return status, msg313 314 315def get_auth_sources(type):316    """Get the authenticated source object from the registry"""317 318    auth_sources = getattr(current_app, '_pgadmin_auth_sources', None)319 320    if auth_sources is None or not isinstance(auth_sources, dict):321        auth_sources = dict()322 323    if type in auth_sources:324        return auth_sources[type]325 326    auth_source = AuthSourceRegistry.get(type)327 328    if auth_source is not None:329        auth_sources[type] = auth_source330        setattr(current_app, '_pgadmin_auth_sources', auth_sources)331 332    return auth_source333 334 335def init_app(app):336    auth_sources = dict()337 338    setattr(app, '_pgadmin_auth_sources', auth_sources)339    AuthSourceRegistry.load_modules(app)340 341    return auth_sources342