Python: Using JWT in cookies with a flask app and restful API!
Tags: python, webdev, security
Published: 2020-01-11 12:36:06

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`](https://en.wikipedia.org/wiki/JSON_Web_Token) 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](https://flask-doc.readthedocs.io/en/latest/) and [flask-restful](https://flask-restful.readthedocs.io/en/latest/).

Need a working example? Here's a [gist](https://gist.github.com/TotallyNotChase/fe0d66318e7f6b56fd2bc0b8972a483c)

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](https://www.grc.com/passwords.htm).
* 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](https://flask-jwt-extended.readthedocs.io/en/stable/api/#utilities)!

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](https://flask-jwt-extended.readthedocs.io/en/stable/api/#module-flask_jwt_extended)!

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](https://flask-jwt-extended.readthedocs.io/en/stable/token_freshness/)!

## 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](https://flask-jwt-extended.readthedocs.io/en/stable/api/#protected-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!