codekingpro/portable-devtools

1"""2    flask_security.recovery_codes3    ~~~~~~~~~~~~~~~~~~~~~~~~4 5    Flask-Security Recovery Codes Module6 7    :copyright: (c) 2022-2024 by J. Christopher Wagner (jwag).8    :license: MIT, see LICENSE for more details.9"""10 11from __future__ import annotations12 13import typing as t14 15from flask import after_this_request, request, redirect16from flask_login import current_user17 18 19from .decorators import anonymous_user_required, auth_required, unauth_csrf20from .forms import (21    build_form_from_request,22    get_form_field_label,23    get_form_field_xlate,24    Form,25    Required,26    StringField,27    SubmitField,28)29from .proxies import _datastore, _security30from .tf_plugin import tf_check_state, tf_illegal_state31from .utils import (32    _,33    base_render_json,34    config_value as cv,35    get_message,36    get_post_login_redirect,37    view_commit,38)39 40if t.TYPE_CHECKING:  # pragma: no cover41    from cryptography.fernet import MultiFernet42    import flask43    from flask.typing import ResponseValue44    from .datastore import User45 46 47class MfRecoveryCodesUtil:48    """Handle creation, checking, encrypting and decrypting recovery codes.49    Since these are rarely used - keep them encrypted until needed - yes50    if someone gets access to memory they can find the key...51    """52 53    def __init__(self, app: flask.Flask):54        self.cryptor: MultiFernet | None = None55        keys = cv("MULTI_FACTOR_RECOVERY_CODES_KEYS", app)56        # N.B. order is important - first key is 'primary'.57        if keys:58            self.setup_cryptor(keys)59 60    def setup_cryptor(self, keys: list[bytes]) -> None:61        from cryptography.fernet import Fernet, MultiFernet62 63        cryptors: list[Fernet] = []64        for key in keys:65            cryptors.append(Fernet(key))66        self.cryptor = MultiFernet(cryptors)67 68    def create_recovery_codes(self, user: User) -> list[str]:69        # Create new recovery codes and store in user record.70        # If configured codes are stored encrypted - but plainttext71        # versions are returned.72        new_codes = _security._totp_factory.generate_recovery_codes(73            cv("MULTI_FACTOR_RECOVERY_CODES_N")74        )75        _datastore.mf_set_recovery_codes(user, self.encrypt_codes(new_codes))76        return new_codes77 78    def get_recovery_codes(self, user: User) -> list[str]:79        ecodes = _datastore.mf_get_recovery_codes(user)80        return self.decrypt_codes(ecodes)81 82    def check_recovery_code(self, user: User, code: str) -> bool:83        # Verify code is valid84        codes = _datastore.mf_get_recovery_codes(user)85        dcodes = self.decrypt_codes(codes)86        return code in dcodes87 88    def delete_recovery_code(self, user: User, code: str) -> bool:89        # codes are single use - so delete after use.90        # encrypting code gives different answer due to time stamp.91        # we don't want to re-encrypt other codes.92        codes = _datastore.mf_get_recovery_codes(user)93        if self.cryptor:94            codes = self.decrypt_codes(codes)95        idx = codes.index(code)96        return _datastore.mf_delete_recovery_code(user, idx)97 98    def encrypt_codes(self, codes: list[str]) -> list[str]:99        if not self.cryptor:100            return codes101        ecodes = []102        for code in codes:103            ecodes.append(self.cryptor.encrypt(code.encode()).decode())104        return ecodes105 106    def decrypt_codes(self, codes: list[str]) -> list[str]:107        from cryptography.fernet import InvalidToken108 109        if not self.cryptor:110            return codes111        dcodes = []112        for code in codes:113            try:114                dcode = self.cryptor.decrypt(115                    code.encode(), cv("MULTI_FACTOR_RECOVERY_CODE_TTL")116                )117                dcodes.append(dcode.decode())118            except InvalidToken:119                # should we delete this?120                pass121        return dcodes122 123 124class MfRecoveryCodesForm(Form):125    """Generate and fetch recovery codes"""126 127    # show_codes is a GET option., generate_new_codes is a POST option128    show_codes = SubmitField(get_form_field_xlate(_("Show Recovery Codes")))129    generate_new_codes = SubmitField(130        get_form_field_xlate(_("Generate New Recovery Codes"))131    )132 133    def __init__(self, *args, **kwargs):134        super().__init__(*args, **kwargs)135 136    def validate(self, **kwargs: t.Any) -> bool:137        if not super().validate(**kwargs):  # pragma: no cover138            return False139        return True140 141 142class MfRecoveryForm(Form):143    """Accept recovery code for second factor authentication"""144 145    code = StringField(146        get_form_field_xlate(_("Recovery Code")),147        validators=[Required()],148    )149    submit = SubmitField(get_form_field_label("submitcode"))150 151    def __init__(self, *args: t.Any, **kwargs: t.Any):152        super().__init__(*args, **kwargs)153        # filled by view154        self.user: User | None = None155 156    def validate(self, **kwargs: t.Any) -> bool:157        if not super().validate(**kwargs):  # pragma: no cover158            return False159        assert self.user is not None160        if not _security._mf_recovery_codes_util.check_recovery_code(161            self.user, self.code.data162        ):163            self.code.errors.append(get_message("INVALID_RECOVERY_CODE")[0])164            return False165        return True166 167 168@auth_required(169    lambda: cv("API_ENABLED_METHODS"),170    within=lambda: cv("FRESHNESS"),171    grace=lambda: cv("FRESHNESS_GRACE_PERIOD"),172)173def mf_recovery_codes() -> ResponseValue:174    """175    Create and download multi-factor recovery codes.176    For forms, we want the user to explicitly request to see the codes - so177    the form has a show_codes submit button.178    """179    form = t.cast(180        MfRecoveryCodesForm, build_form_from_request("mf_recovery_codes_form")181    )182 183    if form.validate_on_submit():184        # generate new codes185        codes = _security._mf_recovery_codes_util.create_recovery_codes(current_user)186        after_this_request(view_commit)187        if _security._want_json(request):188            payload = dict(recovery_codes=codes)189            return base_render_json(form, include_user=False, additional=payload)190        return _security.render_template(191            cv("MULTI_FACTOR_RECOVERY_CODES_TEMPLATE"),192            mf_recovery_codes_form=form,193            recovery_codes=codes,194            **_security._run_ctx_processor("mf_recovery_codes"),195        )196 197    codes = _security._mf_recovery_codes_util.get_recovery_codes(current_user)198    if _security._want_json(request):199        return base_render_json(200            form, include_user=False, additional=dict(recovery_codes=codes)201        )202    show_codes = request.args.get("show_codes", False)203    if show_codes and not codes:204        form.show_codes.errors = []205        form.show_codes.errors.append(get_message("NO_RECOVERY_CODES_SETUP")[0])206    return _security.render_template(207        cv("MULTI_FACTOR_RECOVERY_CODES_TEMPLATE"),208        mf_recovery_codes_form=form,209        recovery_codes=codes if show_codes else [],210        **_security._run_ctx_processor("mf_recovery_codes"),211    )212 213 214@anonymous_user_required215@unauth_csrf()216def mf_recovery():217    """View for entering a recovery code.218 219    User must have already provided valid username/password.220    User must have already established 2FA221 222    """223    form = t.cast(MfRecoveryForm, build_form_from_request("mf_recovery_form"))224    form.user = tf_check_state(["ready"])225    if not form.user:226        return tf_illegal_state(form, cv("TWO_FACTOR_ERROR_VIEW"))227 228    if form.validate_on_submit():229        # Valid code - we want these to be one time - so remove it from list230        _security._mf_recovery_codes_util.delete_recovery_code(231            form.user, form.code.data232        )233        after_this_request(view_commit)234 235        # In the recovery case - don't set/offer validity token.236        _security.two_factor_plugins.tf_complete(form.user, True)237 238        if not _security._want_json(request):239            return redirect(get_post_login_redirect())240        else:241            return base_render_json(form)242 243    if _security._want_json(request):244        return base_render_json(form, include_user=False)245    return _security.render_template(246        cv("MULTI_FACTOR_RECOVERY_TEMPLATE"),247        mf_recovery_form=form,248        **_security._run_ctx_processor("mf_recovery"),249    )250