Search code examples
authenticationflasksessionsession-cookiesflask-login

How do I store one session variable for some time even after closing the browser in Flask?


I am using Flask-Login's function login() with remember=True.

This is working fine. But recently I implemented a Two-factor authentication where the user has to download an authenticator app to be able to login.

I created my own @tfa_required decorator that uses @login_required in case the user enabled TFA. The way I keep track of the login status is with a session variable called auth_level to know whether the user successfully logged in using TFA. (code snippets attached).

Assume that the user logged in to a browser which then closed it. He would still be logged in (Since he is logged in with remember = True), but he will have to enter TFA again (auth_level=1 -- meaning TFA is not successful yet -- now instead of auth_level=2).

What I really want is to store auth_level session variable for as long as Flask-Login remembers that the user logged in, even after closing the browser.

As far as I understand, Flask-Login uses a different session to store the login-related variables so that the user is still logged in for some time.

How do I make the client remember auth_level even after closing the browser?

Thank you.

# Decorator for TFA-required views
def tfa_login_required(f):
    """Implement additional TFA checking, on top of login_required.

    TFA is only checked if it has been activated for the user.
    """
    @wraps(f)
    @login_required
    def decorated_function(*args, **kwargs):
        # if the user is authenticated and has tfa enabled but not tfa passed
        if (
            current_user.tfa_enabled
            and session.get('auth_level') != 2
        ):
            return redirect(url_for('auth.login_view'))
        # if authenticated proceed normally
        else:
            return f(*args, **kwargs)
    return decorated_function

Example of using login_required without tfa_required:

@auth.route('/logout/')
@login_required
def logout():
    logout_user()
    session['auth_level'] = 0
    session.modified = True
    return redirect(url_for('main.home'))

Example using both, tfa_required AND login_required:

@main.route('/some_route/', methods=['GET'])
@tfa_login_required
def some_route():
    do_some_work_that_requires_full_authentication()

Solution

  • I ended up using the Flask-Login cookies themselves to include the variable in the Flask-Login cookie. For reference, here is the code:

    from flask_login import (
    ....
    COOKIE_DURATION,
    COOKIE_SECURE,
    COOKIE_HTTPONLY,
    encode_cookie,
    )
    

    In the view function, instead of returning url_for(...), I create a response object with make_response() and set the response cookie using the following:

        response = make_response(redirect(url_for('...')))
    
        # Setting cookie similar to Flask Login way of setting its cookies
        duration = current_app.config.get('REMEMBER_COOKIE_DURATION', COOKIE_DURATION)
        secure = current_app.config.get('REMEMBER_COOKIE_SECURE', COOKIE_SECURE)
        httponly = current_app.config.get('REMEMBER_COOKIE_HTTPONLY', COOKIE_HTTPONLY)
        expires = duration + datetime.utcnow()
    
        # encode additional data to ensure that the user cannot simply use the value
        # from the remember cookie
        data = encode_cookie(str(session['_user_id']) + '_cookie_variable')
    
        response.set_cookie(
            'cookie_variable',
            value=data,
            expires=expires,
            secure=secure,
            httponly=httponly
        )
        return response