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"""Multi-factor Authentication (MFA) implementation"""10 11from flask import Blueprint, session, Flask12from flask_babel import gettext as _13 14import config15from .utils import mfa_enabled, segregate_valid_and_invalid_mfa_methods16 17from .registry import MultiFactorAuthRegistry18from .views import validate_view, registration_view19 20 21def __create_blueprint() -> Blueprint:22    """23    Geneates the blueprint for 'mfa' endpoint, and also - define the required24    endpoints within that blueprint.25 26    Returns:27        Blueprint: MFA blueprint object28    """29    blueprint = Blueprint(30        "mfa", __name__, url_prefix="/mfa",31        static_folder="static",32        template_folder="templates"33    )34 35    blueprint.add_url_rule(36        "/validate", "validate", validate_view, methods=("GET", "POST",)37    )38 39    blueprint.add_url_rule(40        "/register", "register", registration_view, methods=("GET", "POST",)41    )42 43    return blueprint44 45 46def init_app(app: Flask):47    """48    Initialize the flask application for the multi-faction authentication49    end-points, when the SERVER_MODE is set to True, and MFA_ENABLED is set to50    True in the configuration file.51 52    Args:53        app (Flask): Flask Application object54    """55 56    if getattr(config, "SERVER_MODE", False) is False and \57            getattr(config, "MFA_ENABLED", False) is False:58        return59 60    MultiFactorAuthRegistry.load_modules(app)61 62    def exclude_invalid_mfa_auth_methods():63        """64        Exclude the invalid MFA auth methods specified in MFA_SUPPORTED_METHODS65        configuration.66        """67 68        supported_methods = getattr(config, "MFA_SUPPORTED_METHODS", [])69        invalid_auth_methods = []70 71        supported_methods, invalid_auth_methods = \72            segregate_valid_and_invalid_mfa_methods(supported_methods)73 74        for auth_method in invalid_auth_methods:75            app.logger.warning(_(76                "'{}' is not a valid multi-factor authentication method"77            ).format(auth_method))78 79        config.MFA_SUPPORTED_METHODS = supported_methods80        blueprint = __create_blueprint()81 82        for mfa_method in supported_methods:83            mfa = MultiFactorAuthRegistry.get(mfa_method)84            mfa.register_url_endpoints(blueprint)85 86        app.register_blueprint(blueprint)87        app.register_logout_hook(blueprint)88 89        from flask_login import user_logged_out90 91        @user_logged_out.connect_via(app)92        def clear_session_on_login(sender, user):93            session['mfa_authenticated'] = False94 95    def disable_mfa():96        """97        Set MFA_ENABLED configuration to False.98 99        Also - log a warning message about no valid authentication method found100        during initialization.101        """102        if getattr(config, 'MFA_ENABLED', False) is True and \103                getattr(config, 'SERVER_MODE', False) is True:104            app.logger.warning(_(105                "No valid multi-factor authentication found, hence - "106                "disabling it."107            ))108        config.MFA_ENABLED = False109 110    mfa_enabled(exclude_invalid_mfa_auth_methods, disable_mfa)111