codekingpro/portable-devtools

1"""2    flask_security.tf_plugin3    ~~~~~~~~~~~~~~~~~~~~~~~~4 5    Flask-Security Two-Factor Plugin Module6 7    :copyright: (c) 2022-2024 by J. Christopher Wagner (jwag).8    :license: MIT, see LICENSE for more details.9 10    TODO:11        - add localized callback for select choices.12"""13 14from __future__ import annotations15 16import typing as t17 18from flask import request, redirect, session19 20from .decorators import unauth_csrf21from .forms import (22    build_form_from_request,23    get_form_field_xlate,24    Form,25    RadioField,26    SubmitField,27)28from .proxies import _datastore, _security29from .utils import (30    _,31    base_render_json,32    check_and_get_token_status,33    config_value as cv,34    do_flash,35    get_message,36    get_within_delta,37    get_url,38    login_user,39    propagate_next,40    simple_render_json,41    url_for_security,42)43 44if t.TYPE_CHECKING:  # pragma: no cover45    import flask46    from flask.typing import ResponseValue47    from flask import Response48    from .core import Security49    from .datastore import User50 51 52class TwoFactorSelectForm(Form):53    which = RadioField(get_form_field_xlate(_("Available Second Factor Methods:")))54    submit = SubmitField(get_form_field_xlate(_("Select")))55 56    def __init__(self, *args, **kwargs):57        super().__init__(*args, **kwargs)58 59 60@unauth_csrf()61def tf_select() -> ResponseValue:62    # Ask user which MFA method they want to use.63    # This is used when a user has setup more than one type of 2FA.64    form = t.cast(65        TwoFactorSelectForm, build_form_from_request("two_factor_select_form")66    )67 68    # This endpoint is unauthenticated - make sure we're in a valid state69    if not all(k in session for k in ["tf_user_id", "tf_select"]):70        # illegal call on this endpoint71        tf_clean_session()72        return tf_illegal_state(form, cv("TWO_FACTOR_ERROR_VIEW"))73 74    user = _datastore.find_user(fs_uniquifier=session["tf_user_id"])75    if not user:  # pragma no cover76        # hard to imagine - someone deletes the user while they are logging in.77        tf_clean_session()78        return tf_illegal_state(form, cv("TWO_FACTOR_ERROR_VIEW"))79 80    setup_methods = _security.two_factor_plugins.get_setup_tf_methods(user)81    form.which.choices = setup_methods82 83    if form.validate_on_submit():84        response = None85        tf_impl = _security.two_factor_plugins.method_to_impl(user, form.which.data)86        if tf_impl:87            json_payload = {"tf_required": True}88            response = tf_impl.tf_login(89                user, json_payload, next_loc=propagate_next(request.url, None)90            )91        if not response:  # pragma no cover92            # This really can't happen unless between the time the started logging in93            # and now, they deleted a second factor (which they would have to do94            # in another window).95            tf_clean_session()96            return tf_illegal_state(form, cv("TWO_FACTOR_ERROR_VIEW"))97        return response98 99    if _security._want_json(request):100        payload = {"tf_select": True, "tf_setup_methods": setup_methods}101        return base_render_json(form, include_user=False, additional=payload)102 103    return _security.render_template(104        cv("TWO_FACTOR_SELECT_TEMPLATE"),105        two_factor_select_form=form,106        **_security._run_ctx_processor("tf_select"),107    )108 109 110class TfPluginBase:  # pragma no cover111    def __init__(self, app: flask.Flask):112        pass113 114    def create_blueprint(115        self, app: flask.Flask, bp: flask.Blueprint, state: Security116    ) -> None:117        raise NotImplementedError118 119    def get_setup_methods(self, user: User) -> list[str]:120        """121        Return a list of methods that ``user`` has setup for this second factor122        """123        raise NotImplementedError124 125    def tf_login(126        self, user: User, json_payload: dict[str, t.Any], next_loc: str | None127    ) -> ResponseValue:128        """129        Called from first/primary authenticated views if the user successfully130        authenticated, and required a second method of authentication.131        This method returns the necessary information for the user UI to continue.132        For forms, this is usually a redirect to a secondary sign in form. For JSON133        it is just a payload that describes what the user has to do next.134        """135        raise NotImplementedError136 137 138class TfPlugin:139    """140    Two-Factor plugin support.141 142    Enables multiple independent two-factor implementations to be configured for a given143    app. See TfPluginBase for what a new implementation must provide.144    """145 146    def __init__(self) -> None:147        self._tf_impls: dict[str, TfPluginBase] = {}148 149    def register_tf_impl(150        # N.B. all methods must be unique across all implementations.151        self,152        app: flask.Flask,153        name: str,154        impl: t.Type[TfPluginBase],155    ) -> None:156        self._tf_impls[name] = impl(app)157 158    def create_blueprint(159        self, app: flask.Flask, bp: flask.Blueprint, state: Security160    ) -> None:161        if state.support_mfa:162            for impl in self._tf_impls.values():163                impl.create_blueprint(app, bp, state)164            # Add our route for selecting between multiple active two-factor165            # mechanisms.166            bp.route(167                cv("TWO_FACTOR_SELECT_URL", app),168                methods=["GET", "POST"],169                endpoint="tf_select",170            )(tf_select)171 172    def method_to_impl(self, user: User, method: str) -> TfPluginBase | None:173        # reverse map a method to the implementation.174        # N.B. again - requires that methods be unique across all implementations.175        # There is a small window that a previously setup method was removed.176        for impl in self._tf_impls.values():177            setup_methods = impl.get_setup_methods(user)178            if method in setup_methods:179                return impl180        return None  # pragma no cover181 182    def get_setup_tf_methods(self, user: User) -> list[str]:183        # Return list of methods that user has setup184        methods = []185        for impl in self._tf_impls.values():186            methods.extend(impl.get_setup_methods(user))187        return methods188 189    def tf_enter(190        self,191        user: User,192        remember_me: bool,193        primary_authn_via: str,194        next_loc: str | None,195    ) -> ResponseValue | None:196        """Check if two-factor is required and if so, start the process.197        Must be called in a request context.198        remember_me controls 2 cookies - the remember_me cookie and the tf_validity199        cookie. We use the session to hold the fact that the user requested 'remember'200        across the second factor.201        """202        json_payload: dict[str, t.Any]203        if _security.support_mfa:204            tf_setup_methods = self.get_setup_tf_methods(user)205            if cv("TWO_FACTOR_REQUIRED") or len(tf_setup_methods) > 0:206                tf_fresh = tf_verify_validity_token(user.fs_uniquifier)207                if cv("TWO_FACTOR_ALWAYS_VALIDATE") or not tf_fresh:208                    # Clean out any potential old session info - in case of previous209                    # aborted 2FA attempt.210                    tf_clean_session()211 212                    json_payload = {"tf_required": True}213                    if remember_me:214                        session["tf_remember_login"] = remember_me215 216                    session["tf_user_id"] = user.fs_uniquifier217                    # A backwards compat hack - the original twofactor could be setup218                    # as part of initial login.219                    if len(tf_setup_methods) == 0:220                        # only initial two-factor implementation supports this221                        return self._tf_impls["code"].tf_login(222                            user, json_payload, next_loc223                        )224                    elif len(tf_setup_methods) == 1:225                        # method_to_impl can't return None here since we just226                        # got the methods up above.227                        impl = t.cast(228                            TfPluginBase,229                            self.method_to_impl(user, tf_setup_methods[0]),230                        )231                        return impl.tf_login(user, json_payload, next_loc)232                    else:233                        session["tf_select"] = True234                        if not _security._want_json(request):235                            values = dict(next=next_loc) if next_loc else dict()236                            return redirect(url_for_security("tf_select", **values))237                        # Let's force app to go through tf-select just in case we want238                        # to do further validation... However, provide the choices239                        # so they can just do a POST240                        json_payload.update(241                            {242                                "tf_select": True,243                                "tf_setup_methods": tf_setup_methods,244                            }245                        )246                        return simple_render_json(json_payload)247        return None248 249    def tf_complete(self, user: User, dologin: bool) -> str | None:250        remember = session.pop("tf_remember_login", None)251 252        if dologin:253            login_user(user, remember=remember)254        tf_clean_session()255        token = None256        # return a token to avoid future two-factor prompts (for a period of time)257        if not cv("TWO_FACTOR_ALWAYS_VALIDATE") and remember:258            token = generate_tf_validity_token(user.fs_uniquifier)259        return token260 261 262def generate_tf_validity_token(fs_uniquifier):263    """Generates a unique token for the specified user.264 265    :param fs_uniquifier: The fs_uniquifier of a user to whom the token belongs to266    """267    return _security.tf_validity_serializer.dumps(fs_uniquifier)268 269 270def tf_validity_token_status(token):271    """Returns the expired status, invalid status, and user of a272    Two-Factor Validity token.273    For example::274 275        expired, invalid, user = tf_validity_token_status('...')276 277    :param token: The Two-Factor Validity token278    """279    return check_and_get_token_status(280        token, "tf_validity", get_within_delta("TWO_FACTOR_LOGIN_VALIDITY")281    )282 283 284def tf_verify_validity_token(fs_uniquifier: str) -> bool:285    """Returns the status of the Two-Factor Validity token based on the current286    request.287 288    :param fs_uniquifier: The ``fs_uniquifier`` of the submitting user.289    """290    token = request.cookies.get("tf_validity", default=None)291    if token is None:292        return False293 294    expired, invalid, uniquifier = tf_validity_token_status(token)295    if expired or invalid or (fs_uniquifier != uniquifier):296        return False297 298    return True299 300 301def tf_set_validity_token_cookie(response: Response, token: str) -> Response:302    """Sets the Two-Factor validity token for a specific user given that is303    configured and the user selects remember me304 305    :param response: The response with which to set the set_cookie306    :param token: validity token307    """308    cookie_kwargs = cv("TWO_FACTOR_VALIDITY_COOKIE")309    max_age = int(get_within_delta("TWO_FACTOR_LOGIN_VALIDITY").total_seconds())310    response.set_cookie("tf_validity", value=token, max_age=max_age, **cookie_kwargs)311    # This is likely overkill since so far we only return this on a POST which is312    # unlikely to be cached.313    response.vary.add("Cookie")314    return response315 316 317def tf_check_state(allowed_states: list[str]) -> User | None:318    if (319        not all(k in session for k in ["tf_user_id", "tf_state"])320        or session["tf_state"] not in allowed_states321    ):322        tf_clean_session()323        return None324 325    user = _datastore.find_user(fs_uniquifier=session["tf_user_id"])326    if not user:327        tf_clean_session()328    return user329 330 331def tf_illegal_state(form, redirect_to):332    m, c = get_message("TWO_FACTOR_PERMISSION_DENIED")333    if not _security._want_json(request):334        do_flash(m, c)335        return redirect(get_url(redirect_to))336    else:337        form.form_errors.append(m)338        return base_render_json(form, include_user=False)339 340 341def tf_clean_session():342    """343    Clean out ALL stuff stored in session (e.g. on logout or restart of a session)344    """345    if cv("TWO_FACTOR"):346        for k in [347            "tf_state",348            "tf_user_id",349            "tf_primary_method",350            "tf_remember_login",351            "tf_totp_secret",352            "tf_select",353        ]:354            session.pop(k, None)355