codekingpro/portable-devtools

1"""2    flask_security.two_factor3    ~~~~~~~~~~~~~~~~~~~~~~~~~~~4 5    Flask-Security two_factor module6 7    :copyright: (c) 2016 by Gal Stainfeld, at Emedgene8    :copyright: (c) 2019-2024 by J. Christopher Wagner (jwag).9"""10 11from __future__ import annotations12 13import typing as t14 15from flask import current_app, redirect, request, session16 17from .forms import (18    get_form_field_xlate,19    DummyForm,20    TwoFactorRescueForm,21)22from .proxies import _security, _datastore23from .tf_plugin import TfPluginBase, tf_clean_session24from .utils import (25    _,26    SmsSenderFactory,27    base_render_json,28    config_value as cv,29    do_flash,30    json_error_response,31    send_mail,32    url_for_security,33)34from .signals import (35    tf_code_confirmed,36    tf_disabled,37    tf_security_token_sent,38    tf_profile_changed,39)40 41if t.TYPE_CHECKING:  # pragma: no cover42    import flask43    from .core import Security44    from .datastore import User45    from flask.typing import ResponseValue46 47 48def tf_send_security_token(user, method, totp_secret, phone_number):49    """Sends the security token via email/sms for the specified user.50 51    :param user: The user to send the code to52    :param method: The method in which the code will be sent53                   ('email' or 'sms', or 'authenticator') at the moment54    :param totp_secret: a unique shared secret of the user55    :param phone_number: If 'sms' phone number to send to56 57    There is no return value - it is assumed that exceptions are thrown by underlying58    methods that callers can catch.59 60    Flask-Security code should NOT call this directly -61    call :meth:`.UserMixin.tf_send_security_token`62    """63    token_to_be_sent = _security._totp_factory.generate_totp_password(totp_secret)64    if method == "email" or method == "mail":65        send_mail(66            cv("EMAIL_SUBJECT_TWO_FACTOR"),67            user.email,68            "two_factor_instructions",69            user=user,70            token=token_to_be_sent,71            username=user.calc_username(),72        )73    elif method == "sms":74        msg = f"Use this code to log in: {token_to_be_sent}"75        from_number = cv("SMS_SERVICE_CONFIG")["PHONE_NUMBER"]76        to_number = phone_number77        sms_sender = SmsSenderFactory.createSender(cv("SMS_SERVICE"))78        sms_sender.send_sms(from_number=from_number, to_number=to_number, msg=msg)79 80    else:81        # password are generated automatically in the authenticator apps or not needed82        token_to_be_sent = None83 84    tf_security_token_sent.send(85        current_app._get_current_object(),86        _async_wrapper=current_app.ensure_sync,87        user=user,88        method=method,89        token=token_to_be_sent,90        login_token=token_to_be_sent,91        phone_number=phone_number,92    )93 94 95def complete_two_factor_process(user, primary_method, totp_secret, is_changing):96    """clean session according to process (login or changing two-factor method)97    and perform action accordingly98    """99 100    _datastore.tf_set(user, primary_method, totp_secret=totp_secret)101 102    # if we are changing two-factor method103    dologin = False104    if is_changing:105        completion_message = "TWO_FACTOR_CHANGE_METHOD_SUCCESSFUL"106        tf_profile_changed.send(107            current_app._get_current_object(),108            _async_wrapper=current_app.ensure_sync,109            user=user,110            method=primary_method,111        )112    # if we are logging in for the first time113    else:114        completion_message = "TWO_FACTOR_LOGIN_SUCCESSFUL"115        tf_code_confirmed.send(116            current_app._get_current_object(),117            _async_wrapper=current_app.ensure_sync,118            user=user,119            method=primary_method,120        )121        dologin = True122    token = _security.two_factor_plugins.tf_complete(user, dologin)123    return completion_message, token124 125 126def set_rescue_options(form: TwoFactorRescueForm, user: User) -> dict[str, str]:127    # Based on config - set up options for rescue.128    # Note that this modifies the passed in Form as well as returns129    # a dict that can be returned as part of a JSON response.130    recovery_options = dict(help=url_for_security("two_factor_rescue"))131 132    if cv("TWO_FACTOR_RESCUE_EMAIL"):133        recovery_options["email"] = url_for_security("two_factor_rescue")134        form.help_setup.choices.append(135            ("email", get_form_field_xlate(_("Send code via email")))136        )137 138    if (139        _security.support_mfa140        and cv("MULTI_FACTOR_RECOVERY_CODES")141        and _datastore.mf_get_recovery_codes(user)142    ):143        recovery_options["recovery_code"] = url_for_security("mf_recovery")144        form.help_setup.choices.append(145            (146                "recovery_code",147                get_form_field_xlate(_("Use previously downloaded recovery code")),148            )149        )150    return recovery_options151 152 153def tf_disable(user):154    """Disable two factor for user"""155    tf_clean_session()156    _datastore.tf_reset(user)157    tf_disabled.send(158        current_app._get_current_object(),159        _async_wrapper=current_app.ensure_sync,160        user=user,161    )162 163 164def is_tf_setup(user):165    """Return True is user account is setup for 2FA."""166    return user.tf_totp_secret and user.tf_primary_method167 168 169class CodeTfPlugin(TfPluginBase):170    def __init__(self, app: flask.Flask):171        super().__init__(app)172 173    def create_blueprint(174        self, app: flask.Flask, bp: flask.Blueprint, state: Security175    ) -> None:176        pass177 178    def get_setup_methods(self, user: User) -> list[str]:179        if is_tf_setup(user):180            assert user.tf_primary_method is not None181            return [user.tf_primary_method]182        return []183 184    def tf_login(185        self, user: User, json_payload: dict[str, t.Any], next_loc: str | None186    ) -> ResponseValue:187        """Helper for two-factor authentication login188 189        This is called only when login/password have already been validated.190        This can be from login, register, confirm, unified sign in, unified magic link.191 192        If two-factor is already setup then this sends a code if the method requires it.193        If not, then user is redirected to two-factor-setup.194        In either case we do NOT log in user, so we must store some info in session to195        track our state (including what user).196        """197 198        # if user's two-factor properties are not configured199        if not is_tf_setup(user):200            session["tf_state"] = "setup_from_login"201            json_payload["tf_state"] = "setup_from_login"202            if not _security._want_json(request):203                return redirect(url_for_security("two_factor_setup"))204 205        # if user's two-factor properties are configured206        else:207            session["tf_state"] = "ready"208            json_payload["tf_state"] = "ready"209            json_payload["tf_primary_method"] = user.tf_primary_method210            json_payload["tf_method"] = user.tf_primary_method211 212            if user.tf_primary_method in ["mail", "email", "sms"]:213                msg = user.tf_send_security_token(214                    method=user.tf_primary_method,215                    totp_secret=user.tf_totp_secret,216                    phone_number=getattr(user, "tf_phone_number", None),217                )218                if msg:219                    # send code didn't work220                    if not _security._want_json(request):221                        # This is a mess -222                        # we are deep down in the login/unified sign in flow.223                        do_flash(msg, "error")224                        return redirect(url_for_security("login"))225                    else:226                        payload = json_error_response(errors=msg)227                        return _security._render_json(payload, 500, None, None)228 229            if not _security._want_json(request):230                values = dict(next=next_loc) if next_loc else dict()231                return redirect(232                    url_for_security("two_factor_token_validation", **values)233                )234 235        # JSON response - Fake up a form - doesn't really matter which.236        form = DummyForm(formdata=None)237        return base_render_json(form, include_user=False, additional=json_payload)238