from flask.sessions import SecureCookieSessionInterface
from flask import g
from flask_login import user_loaded_from_request


@user_loaded_from_request.connect
def user_loaded_from_request(app, user=None):
    g.login_via_request = True


class JWTSessionInterface(SecureCookieSessionInterface):
    """This is a custom session interface that will turn off cookies for flask requests so that JWT tokens can be used instead."""
    def should_set_cookie(self, app, session):
        return False

    def save_session(self, app, session, response):
        if g.get('login_via_request'):
            return
        return super(JWTSessionInterface, self).save_session(app, session, response)