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) utility functions"""10 11from collections.abc import Callable12from functools import wraps13 14from flask import url_for, session, request, redirect15from flask_login.utils import login_url16from flask_security import current_user, login_required17 18import config19from pgadmin.model import UserMFA, db20from .registry import MultiFactorAuthRegistry21 22 23class ValidationException(Exception):24    """25    class: ValidationException26    Base class: Exception27 28    An exception class for raising validation issue.29    """30    pass31 32 33def segregate_valid_and_invalid_mfa_methods(34    mfa_supported_methods: list35) -> (list, list):36    """37    Segregate the valid and invalid authentication methods from the given38    methods.39 40    Args:41        mfa_supported_methods (list): List of auth methods42 43    Returns:44        list, list: Set of valid & invalid auth methods45    """46 47    invalid_auth_methods = []48    valid_auth_methods = []49 50    for mfa in mfa_supported_methods:51 52        # Put invalid MFA method in separate list53        if mfa not in MultiFactorAuthRegistry._registry:54            if mfa not in invalid_auth_methods:55                invalid_auth_methods.append(mfa)56            continue57 58        # Exclude the duplicate entries59        if mfa in valid_auth_methods:60            continue61 62        valid_auth_methods.append(mfa)63 64    return valid_auth_methods, invalid_auth_methods65 66 67def mfa_suppored_methods() -> dict:68    """69    Returns the dictionary containing information on all supported methods with70    information about whether they're registered for the current user, or not.71 72    It returns information in this format:73    {74        <auth_method_name>: {75            "mfa": <MFA Auth Object>,76            "registered": True|False77        },78        ...79    }80 81    Returns:82        dict: List of all supported MFA methods with the flag for the83              registered with the current user or not.84    """85    supported_mfa_auth_methods = dict()86 87    for auth_method in config.MFA_SUPPORTED_METHODS:88        registry = MultiFactorAuthRegistry.get(auth_method)89        supported_mfa_auth_methods[registry.name] = {90            "mfa": registry, "registered": False91        }92 93    auths = UserMFA.query.filter_by(user_id=current_user.id).all()94 95    for auth in auths:96        if auth.mfa_auth in supported_mfa_auth_methods:97            supported_mfa_auth_methods[auth.mfa_auth]['registered'] = True98 99    return supported_mfa_auth_methods100 101 102def user_supported_mfa_methods():103    """104    Returns the dict for the authentication methods, registered for the105    current user, among the list of supported.106 107    Returns:108        dict: dict for the auth methods109    """110    auths = UserMFA.query.filter_by(user_id=current_user.id).all()111    res = dict()112    supported_mfa_auth_methods = dict()113 114    if len(auths) > 0:115        for auth_method in config.MFA_SUPPORTED_METHODS:116            registry = MultiFactorAuthRegistry.get(auth_method)117            supported_mfa_auth_methods[registry.name] = registry118 119        for auth in auths:120            if auth.mfa_auth in supported_mfa_auth_methods:121                res[auth.mfa_auth] = \122                    supported_mfa_auth_methods[auth.mfa_auth]123 124    return res125 126 127def is_mfa_session_authenticated() -> bool:128    """129    Checks if this session is authenticated, or not.130 131    Returns:132        bool: Is this session authenticated?133    """134    return session.get('mfa_authenticated', False) is True135 136 137def mfa_enabled(execute_if_enabled, execute_if_disabled) -> None:138    """139    A ternary method to enable calling either of the methods based on the140    configuration for the MFA.141 142    When MFA is enabled and has a valid supported auth methods,143    'execute_if_enabled' method is executed, otherwise -144    'execute_if_disabled' method is executed.145 146    Args:147        execute_if_enabled (Callable[[], None]):  Method to executed when MFA148                                                  is enabled.149        execute_if_disabled (Callable[[], None]): Method to be executed when150                                                  MFA is disabled.151 152    Returns:153        None: Expecting the methods to return None as it will not be consumed.154 155    NOTE: Removed the typing anotation as it was giving errors.156    """157 158    is_server_mode = getattr(config, 'SERVER_MODE', False)159    enabled = getattr(config, "MFA_ENABLED", False)160    supported_methods = getattr(config, "MFA_SUPPORTED_METHODS", [])161 162    if is_server_mode is True and enabled is True and \163            isinstance(supported_methods, list):164        supported_methods, _ = segregate_valid_and_invalid_mfa_methods(165            supported_methods166        )167 168        if len(supported_methods) > 0:169            return execute_if_enabled()170 171    return execute_if_disabled()172 173 174def mfa_user_force_registration_required(register, not_register) -> None:175    """176    A ternary method to cenable calling either of the methods based on the177    condition force registration is required.178 179    When force registration is enabled, and the current user has not registered180    for any of the supported authentication method, then the 'register' method181    is executed, otherwise - 'not_register' method is executed.182 183    Args:184        register (Callable[[], None])    : Method to be executed when for185                                           registration required and user has186                                           not registered for any auth method.187        not_register (Callable[[], None]): Method to be executed otherwise.188 189    Returns:190        None: Expecting the methods to return None as it will not be consumed.191    """192    return register() \193        if getattr(config, "MFA_FORCE_REGISTRATION", False) is True else \194        not_register()195 196 197def mfa_user_registered(registered, not_registered) -> None:198    """199    A ternary method to enable calling either of the methods based on the200    condition - if the user is registed for any of the auth methods.201 202    When current user is registered for any of the supported auth method, then203    the 'registered' method is executed, otherwise - 'not_registered' method is204    executed.205 206    Args:207        registered (Callable[[], None])    : Method to be executed when208                                             registered.209        not_registered (Callable[[], None]): Method to be executed when not210                                             registered211 212    Returns:213        None: Expecting the methods to return None as it will not be consumed.214 215    NOTE: Removed the typing anotation as it was giving errors.216    """217 218    return registered() if len(user_supported_mfa_methods()) > 0 else \219        not_registered()220 221 222def mfa_session_authenticated(authenticated, unauthenticated):223    """224    A ternary method to enable calling either of the methods based on the225    condition - if the user has already authenticated, or not.226 227    When current user is already authenticated, then 'authenticated' method is228    executed, otherwise - 'unauthenticated' method is executed.229 230    Args:231        authenticated (Callable[[], None])  : Method to be executed when232                                              user is authenticated.233        unauthenticated (Callable[[], None]): Method to be executed when the234                                              user is not passed the235                                              authentication.236 237    Returns:238        None: Expecting the methods to return None as it will not be consumed.239 240    NOTE: Removed the typing anotation as it was giving errors.241    """242    return authenticated() if session.get('mfa_authenticated', False) is True \243        else unauthenticated()244 245 246def mfa_required(wrapped):247    """248    A decorator do decide the next course of action when a page is being249    opened, it will open the appropriate page in case the 2FA is not passed.250 251    Function executed252        |253    Check for MFA Enabled? --------+254        |                          |255        | No                       |256        |                          | Yes257    Run the wrapped function [END] |258                                   |259        Is user has registered for at least one MFA method? -+260                    |                                        |261                    | No                                     |262                    |                                        |263        Is force registration required? -+                   |264                    |                    |                   | Yes265                    | No                 |                   |266                    |                    | Yes               |267        Run the wrapped function [END]   |                   |268                                         |                   |269                            Open Registration page [END]     |270                                                              |271                                          Open the authentication page [END]272 273    Args:274        func(Callable[..]): Method to be called if authentcation is passed275    """276 277    def get_next_url():278        next_url = request.url279        registration_url = url_for('mfa.register')280 281        if next_url.startswith(registration_url):282            return url_for('browser.index')283 284        return next_url285 286    def redirect_to_mfa_validate_url():287        return redirect(login_url("mfa.validate", next_url=get_next_url()))288 289    def redirect_to_mfa_registration():290        return redirect(login_url("mfa.register", next_url=get_next_url()))291 292    @wraps(wrapped)293    @login_required294    def inner(*args, **kwargs):295        def execute_func():296            session['mfa_authenticated'] = True297            return wrapped(*args, **kwargs)298 299        def if_else_func(_func, first, second):300            def if_else_func_inner():301                return _func(first, second)302            return if_else_func_inner303 304        return mfa_enabled(305            if_else_func(306                mfa_session_authenticated,307                execute_func,308                if_else_func(309                    mfa_user_registered,310                    redirect_to_mfa_validate_url,311                    if_else_func(312                        mfa_user_force_registration_required,313                        redirect_to_mfa_registration,314                        execute_func315                    )316                )317            ),318            execute_func319        )320 321    return inner322 323 324def is_mfa_enabled() -> bool:325    """326    Returns True if MFA is enabled otherwise False327 328    Returns:329        bool: Is MFA Enabled?330    """331    return mfa_enabled(lambda: True, lambda: False)332 333 334def mfa_delete(auth_name: str) -> bool:335    """336    A utility function to delete the auth method for the current user from the337    configuration database.338 339    Args:340        auth_name (str): Name of the argument341 342    Returns:343        bool: True if auth method was registered for the current user, and344              delete successfully, otherwise - False345    """346    auth = UserMFA.query.filter_by(347        user_id=current_user.id, mfa_auth=auth_name348    )349 350    if int(auth.count()) != 0:351        auth.delete()352        db.session.commit()353 354        return True355 356    return False357 358 359def mfa_add(auth_name: str, options: str) -> None:360    """361    A utility funtion to add/update the auth method in the configuration362    database for the current user with the method specific options.363 364    e.g. email-address for 'email' method, and 'secret' for the 'authenticator'365 366    Args:367        auth_name (str): Name of the auth method368        options (str)  : A data options specific to the auth method369    """370    auth = UserMFA.query.filter_by(371        user_id=current_user.id, mfa_auth=auth_name372    ).first()373 374    if auth is None:375        auth = UserMFA(376            user_id=current_user.id,377            mfa_auth=auth_name,378            options=options379        )380        db.session.add(auth)381 382    # We will override the existing options383    auth.options = options384 385    db.session.commit()386 387 388def fetch_auth_option(auth_name: str) -> (str, bool):389    """390    A utility function to fetch the extra data, stored as options, for the391    given auth method for the current user.392 393    Returns a set as (data, Auth method registered?)394 395    Args:396        auth_name (str): Name of the auth method397 398    Returns:399        (str, bool): (data, has current user registered for the auth method?)400    """401    auth = UserMFA.query.filter_by(402        user_id=current_user.id, mfa_auth=auth_name403    ).first()404 405    if auth is None:406        return None, False407 408    return auth.options, True409