Python: Using JWT in cookies with a flask app and restful API!

This guide aims to provide an in-depth tutorial on how to set up flask-jwt-extended using cookies. There's a LOT of docs online but they are mostly using authentication headers and a frontend framework like react. I wanted to share my experience with using jwt through just the backend (e.g returning redirects instead of json objs). This will also outline some of the guidelines about token freshness, access/refresh token expiry and refreshing expired tokens (through the backend!!).

Note: Usually the backend only returns a jsonify-ed dict with the tokens (when using headers) or simply a status report (when using cookies). Once the frontend framework receives the tokens/status report, it redirects to the desired url. We're just achieving all that without the frontend.

I'll be assuming you know what JWT is and are aware of the python package flask-jwt-extended. You don't need to know how to use the package, I'll show how to do that in the guide :)
You'll also need to know the basics of flask and flask-restful.

Need a working example? Here's a gist

Alright, time to make the flask app itself!

Initializing our app

First let's see what imports we'll be needing. All of these will be explained as we go on, so don't worry too much about them. Just know that these are the functions we are using later if you get confused and ask "wait what does that function do and where did it come from?"

from flask import Flask, redirect, make_response, render_template
from flask_restful import Api
from flask_jwt_extended import (JWTManager, jwt_required, 
                                jwt_refresh_token_required, 
                                jwt_optional, fresh_jwt_required, 
                                get_raw_jwt, get_jwt_identity,
                                create_access_token, create_refresh_token, 
                                set_access_cookies, set_refresh_cookies, 
                                unset_jwt_cookies,unset_access_cookies)

For the purposes of this guide, we're declaring all functions in the same file __init__.py. NEVER DO THIS, always separate your JWT handlers on a different file.

let's make a __init__.py and initialize all the required stuff

app = Flask(__name__)
app.config['BASE_URL'] = 'http://127.0.0.1:5000'  #Running on localhost
app.config['JWT_SECRET_KEY'] = 'super-secret'  # Change this!
app.config['JWT_TOKEN_LOCATION'] = ['cookies']
app.config['JWT_COOKIE_CSRF_PROTECT'] = True
app.config['JWT_CSRF_CHECK_FORM'] = True
jwt = JWTManager(app) 

Don't worry too much about the imports yet, in the above snipper we are only really concerned about Flask, redirect, make_response, and JWTManager. All we're doing here is creating a Flask app instance and then setting up the JWTManager.

As for the configs:-

• We need a JWT_SECRET_KEY for the JWT encryption. Try to go for a completely random string of alphanumeric characters. You can generate them using this.
• We have to set JWT_TOKEN_LOCATION to ['cookies'] if we want to use cookies as our preferred token storage. The default is always auth header.
• We also set JWT_COOKIE_CSRF_PROTECT to True for cookie security, this is True by default so you don't have to do this.
• Lastly, we set JWT_CSRF_CHECK_FORM to True. This will come in handy when we have to pass tokens to api resources. Don't worry about it much yet.

Setting cookies and returning redirects

Let's make a function to create and assign one access token and one refresh token to the cookies. Remember, the tokens are always assigned to a flask response object which must then be returned.

A response object can be anything from a jsonify-ed dict to a redirect. There's a LOT of docs on returning jsonify-ed dicts so I'll be showing how to return a redirect instead.

def assign_access_refresh_tokens(user_id, url):
    access_token = create_access_token(identity=str(user_id))
    refresh_token = create_refresh_token(identity=str(user_id))
    resp = make_response(redirect(url, 302))
    set_access_cookies(resp, access_token)
    set_refresh_cookies(resp, refresh_token)
    return resp

We take in the url to redirect to as an argument. We also take a unique identifer (user_id), which is essential to create a token

Explanation : create_access_token(identity) and create_refresh_token(identity) will create and return an access and refresh token based on the identity provided. Remember the identity should be unique for each user. We then have to make a flask response object that will redirect us to the desired url. We need to attach the cookies to this response object.
make_response(redirect(url, 302)) will create and return a response object that redirects to url with code 302.
Now we just need to assign the access and refresh tokens. We do that using set_access_cookies(response, access_token) and set_refresh_cookies(response, refresh_token).
Finally we return the response object that now has the cookies attached. This will redirect the user to the give url and the user session now has the desired tokens, Awesome!

Read more about jwt functions!

So, that's how you assign tokens, how do you un-assign them? You'd ideally need to revoke those tokens if a user logs out so let's see how to do that!

def unset_jwt():
    resp = make_response(redirect(app.config['BASE_URL'] + '/', 302))
    unset_jwt_cookies(resp)
    return resp

Explanation : yep, it's that easy. You just use unset_jwt_cookies with a response object and return that response. Now the user session no longer has access OR refresh tokens.

Note : This removes both access and refresh tokens. To remove only one of them, use unset_access_cookies(response) or unset_refresh_cookies(response) instead.

Now we have to use jwt decorator functions to handle operations when the user doesn't have tokens/ has invalid tokens/ expired tokens. Let's do that!

@jwt.unauthorized_loader
def unauthorized_callback(callback):
    # No auth header
    return redirect(app.config['BASE_URL'] + '/signup', 302)

@jwt.invalid_token_loader
def invalid_token_callback(callback):
    # Invalid Fresh/Non-Fresh Access token in auth header
    resp = make_response(redirect(app.config['BASE_URL'] + '/signup'))
    unset_jwt_cookies(resp)
    return resp, 302

@jwt.expired_token_loader
def expired_token_callback(callback):
    # Expired auth header
    resp = make_response(redirect(app.config['BASE_URL'] + '/token/refresh'))
    unset_access_cookies(resp)
    return resp, 302

Explanation : Let's discuss what each decorator does first.

• @jwt.unauthorized_loader is called when the user has no access tokens in their request.
• @jwt.invalid_token_loader is called when the user has an invalid access token in their request. This happens when someone tries to access an endpoint using forged tokens.
• @jwt.expired_token_loader is called when the user has an expired but otherwise valid access token in their request.

Read more about jwt loaders!

Now you may have noticed, tokens expire and when it does we unset said token and redirect the user to http://127.0.0.1:5000/token/refresh. What does this url do? Well, it refreshes the expired token if the user has a valid refresh token.

@app.route('/token/refresh', methods=['GET'])
@jwt_refresh_token_required
def refresh():
    # Refreshing expired Access token
    user_id = get_jwt_identity()
    access_token = create_access_token(identity=str(user_id))
    resp = make_response(redirect(app.config['BASE_URL'] + '/', 302))
    set_access_cookies(resp, access_token)
    return resp

Explanation : This function checks whether the user has a valid refresh token first using @jwt_refresh_token_required, if they do, it'll extract the identity given to the token during it's creation using get_jwt_identity() and use it to create another access token. We then assign and return it just like before!

Ok! That's all the jwt handlers, but before we put them to the test. Let's discuss on an important aspect of access token security - token freshness.

Token Freshness

From the docs

[Fresh pattern] is useful for allowing fresh tokens to do some critical things (such as update an email address or complete an online purchase), but to deny those features to non-fresh tokens.

Let's think of a situation, as long the refresh token is valid and non-expired, we can keep generating access tokens, right? If an access token goes into the wrong hands, the attacker will only have a short period of time to use them as expire quite fast (15 mins by default). But what if the attacker gets their hands on a refresh token, which usually have a longer expiry (1 day by default)?

This is why we use the fresh/non-fresh pattern. When creating access tokens using create_access_token we can set fresh = True to create a fresh access token. By default when an access token is created without a fresh value supplied, it will be non-fresh.

We can create a fresh access token by create_access_token(identity = id, fresh = True).

Now we can only let refresh tokens create non-fresh access tokens. non-fresh tokens cannot access endpoints protected with @fresh_jwt_required, only @jwt_required. Fresh tokens can access endpoints protected with both @fresh_jwt_required and @jwt_required. Make sure fresh tokens are only created when the user actually logs in using their password.

Read!

Usage

Now we can finally put the jwt handlers to the test! To use these handlers we have to decorate our functions using jwt required decorators. Here's an overview on them!

• jwt_required will block access to everyone unless they have a valid, non-expired non-fresh or fresh access token.
• fresh_jwt_required will block access to everyone unless they have a valid, non-expired fresh access token.
• jwt_optional won't block access but it'll still allow the jwt to be used inside the function if it is present. Otherwise get_jwt_identity() will return None.
• @jwt_refresh_token_required will block access to everyone unless they have a valid non-expired refresh token

Read more about jwt endpoint decorators!

Let us first assign the tokens first, we'll ideally want to assign them when the user signs up or logs in. Like so

@app.route('/login', methods=['POST'])
def login():
    // Verify username and password //
    return assign_access_refresh_tokens(username , app.config['BASE_URL'] + '/')

Great now we have tokens in the session! Let's declare some functions using the decorators!

@app.route('/account')
@fresh_jwt_required
def account():
     username = get_jwt_identity()
    // very important account settings //

@app.route('/service', methods=['GET'])
@jwt_required
def services():
    username = get_jwt_identity()
    // Not important stuff but still needs to be logged in //

@app.route('/')
@jwt_optional
def index():
    username = get_jwt_identity()   # None if not logged in
    // Accessible to everyone but maybe different to logged in users //

The first endpoint, account can only be accessed by a fresh jwt. Someone who recently logged in will be able to access this as long as their fresh jwt does not expire. When it does expire, the refresh token won't create a fresh jwt and hence to access this, the user will have to log in again.

The second endpoint, services can be accessed by either a fresh or non-fresh jwt. A logged in user can access this for the entirety of their refresh token lifetime without logging in again.

The third endpoint, index can be accessed by anyone. However it may offer a different look/functionality if the user is logged in, i.e get_jwt_identity is not None.

Once the user is done looking around they'd wanna logout, we need to unset their cookies if they do so.

@app.route('/logout')
@jwt_required
def logout():
    # Revoke Fresh/Non-fresh Access and Refresh tokens
    return unset_jwt(), 302

That's it! Now let's learn how to pass the tokens correctly into api resources

Passing JWT to RESTful API resources

Let's assume we have a Class in another module like so

from flask_restful import Resource
from flask_jwt_extended import fresh_jwt_required, get_jwt_identity
class AccountSettings(Resource):
    @fresh_jwt_required
    def post(self):
        username = get_jwt_identity()
        // important setting change commits //

and we add this to our api in __init__.py using api.add_resource(AccountSettings, '/account/change_settings')

So now whenever the url http://127.0.0.1:5000/account/change_settings receives a POST, it'll execute the post function in AccountSettings. However a POST function in an api resource` **cannot* access the CSRF token in our JWT cookies directly. Without that token, the jwt is useless. The reason this happens, along with it's solution is written in the docs-

Typically JWT is used with API servers using JSON payloads, often via AJAX. However you may have an endpoint that receives POST requests directly from an HTML form. Without AJAX, you can’t set the CSRF headers to pass your token to the server. In this scenario you can send the token in a hidden form field. To accomplish this, first configure JWT to check the form for CSRF tokens. Now it’s not necessary to send the csrf in a separate cookie, you can render it directly into your HTML template.

All we have to do is turn the JWT_CSRF_CHECK_FORM option on (which we already did earlier) and put the token in a hidden input field. You can do this by putting <input hidden name="csrf_token" value="{{ csrf_token }}"> INSIDE your form that will POST to the api resource. Now you can use render_template to pass the csrf_token to the template.

Get the csrf_token using csrf_token = (get_raw_jwt() or {}).get("csrf") and pass it to the DOM like so
return render_template("account.html", csrf_token=csrf_token), 200

Note : You can only extract the csrf_token (for access tokens) in endpoints decorated with @fresh_jwt_required or @jwt_required. @jwt_optional will also work however, if there is no JWT present, it'll return None.

Note : You only have to do this on RESTful API resources that use POST methods. GET method can access the csrf protected jwt just fine.

JWT Error handling in API resources

You may notice the jwt loaders (e.g @jwt.unauthorized_loader, @jwt.invalid_token_loader) not getting hit when the API resource URLs are directly accessed. This is because of flask's default exception handling behaviour, to fix this, simply set PROPAGATE_EXCEPTIONS to True.
app.config['PROPAGATE_EXCEPTIONS'] = True

Advices, Standards and Discussions

• Expiry : The default expiry for access and refresh tokens is 15 minutes and 1 day respectively. You can change this using app.config['JWT_ACCESS_TOKEN_EXPIRES'] or app.config['JWT_REFRESH_TOKEN_EXPIRES'] and assigning a datetime.timedelta() value. It is recommended to keep the access token duration low as it can simply be refreshed.
• Cookie path : It's always nice to only send the tokens to the routes where they are needed, instead of sending them to every route by default. If I wanted to send the access tokens to only /accounts and /services I'd use app.config['JWT_ACCESS_COOKIE_PATH'] = ['/account', '/services']. As you can see it's simply a list of the urls you want the access tokens to be accessible from.

As for refresh tokens, I always send them only to the refresh url.
app.config['JWT_REFRESH_COOKIE_PATH'] = '/token/refresh'

• Verify jwt identity : Even if get_jwt_identity returns an actual id, check whether the id exists in your database before assuming them to be an user outright.
• Fresh/Non-Fresh pattern : Make sure the important routes are protected using fresh tokens and never let fresh tokens be refreshed using refresh tokens. They should only be created when the user actually uses their password/2FA to log in.
• Change up the tokens on password change : It's a good idea to recreate both access and refresh tokens if a user changes their password.

And That's it! I hope this guide helps you in your JWT adventures!