Building REST APIs With Python and Flask: Part 3

This is the third part of a three-part tutorial on creating REST APIs with Python and Flask. In this part, you will learn how to add authentication to REST APIs using JWT tokens.

Here are the links to other parts of this tutorial.

Part 1
Part 2
Part 3(You are here)

JWT Token Authentication

You may have noticed a problem with our API. Anyone can do any operation with our API. There is no access control. We are going to fix that next.

We will use JWT tokens for authenticating users. There is a python package, PyJWT, that can help us to generate tokens. We need to install it first using pipenv.

pipenv install pyjwt

For signing the tokens, we need a secret value. Open our config.py file and add the following line to it.

JWT_SECRET = "super-secret-string"

Now let us start writing our authentication blueprint. Create a new auth.py file and put the following code in it.

import jwt
from flask import Blueprint, current_app, request
from werkzeug.security import check_password_hash

from .db import get_db

blueprint = Blueprint("auth", __name__, url_prefix="/auth")


def generate_token(user_id):
    payload = {"id": user_id}
    return jwt.encode(payload, current_app.config["JWT_SECRET"], algorithm="HS256")


def decode_token(token):
    try:
        payload = jwt.decode(
            token, current_app.config["JWT_SECRET"], algorithms=["HS256"]
        )
    except jwt.exceptions.InvalidSignatureError:
        return None

    return payload["id"]

There are few unused imports. We will use those in the coming sections.
We created the auth blueprint with the url prefix /auth. Refer to the previous tutorials if you do not remember how blueprints work.
We have a generate_token function that accepts a user_id as its only parameter. We then create a dictionary with this user id. This will be our payload for creating tokens. Then we call the encode function from jwt package and pass our payload, the secret and the algorithm to use. HS256 is strong enough for our use cases. You can find some links at the end of this chapter if you want to learn more about JWT and its algorithms.
We have another function decode_token that takes a token generated using generate_token function. We use the decode method from jwt package to recreate payload form token. Note that here the keyword argument name is algorithms and it is a list.
If the token is invalid, jwt will throw an exception. If that happens, we return None.
If everything goes well, we take out the id from the payload and return it.

Next, we need an API that accepts an email and a password. If both email and password are correct, this API will create a JWT token for that user and returns it.

Add the following code to the end of our auth.py file.

@blueprint.route("/get-token/", methods=["POST"])
def get_token():
    data = request.json
    email = data.get("email")
    password = data.get("password")

    db = get_db()
    cursor = db.execute("SELECT * FROM user where email=?", (email,))
    user = cursor.fetchone()

    if user is None:
        return {"error": "Incorrect email"}, 401

    if not check_password_hash(user["password"], password):
        return {"error": "Incorrect password"}, 401

    token = generate_token(user["id"])

    return {"user_id": user["id"], "token": token}

The authentication API uses POST method.
We extract the email and password from the JSON data user submitted.
We then try to select the user with this email id from our database. If there is no user with that email id, we return an error message with a status 401. This status code is generally used for authentication errors.
Passwords stored in our database is in hashed and salted format. We cannot directly compare the value in our database with the value user submitted. For that, we use check_password_hash from werkzeug. The order of the arguments is important here. The first argument is the hashed password and the second is the plaintext password. If check_password_hash returns false, we return an error to the client.
If both email and password match a user, we generate a token using the function we added in the last section. We then return this token and the user’s id to the client.

Our auth API is done. Let us test this to make sure it is working as expected. In postman set the URL to http://127.0.0.1:5000/auth/get-token/ and method to POST. Set the request body type as JSON and put a valid email and password in the data.

{
  "email": "harry@example.com",
  "password": "1234"
}

Hit send and you should see a response with content like this.

{
  "token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpZCI6MX0.aLguoZFXtq0AEUrmAm4CrCUKKq9zt4h8ZmDE5QWgSdw",
  "user_id": 1
}

Try with different combinations of correct and incorrect emails and passwords. See what happens.

Restricting Access

In the last section, we generated authentication tokens for users. In this chapter, we will use this token to restrict access to our API.

When someone uses our API, we expect them to include the token in the request header. If there is a token in the header, we can then decode the token and get the user id. With the user id, we can fetch the user details from the database. Doing this in every API is tedious and cause code duplication. Instead, we will add a function to do this and instruct flask to call this function whenever a request comes in.

We will add this function to our auth.py file since this is related to authentication. First, update the flask import statement to include g.

from flask import Blueprint, current_app, g, request

Then add the following code to the end of the file.

@blueprint.before_app_request
def set_user():
    token = request.headers.get("Authorization")

    if not token:
        g.user = None
        return

    user_id = decode_token(token)

    if not user_id:
        g.user = None
        return

    db = get_db()
    cursor = db.execute("SELECT * FROM user WHERE id=?", (user_id,))
    user = cursor.fetchone()

    g.user = user

blueprint.before_app_request decorator tells flask this function should be called every time a new request comes in. Flask will call the set_user function even if the request is not for an API in auth blueprint.
request.headers is a dictionary containing headers sent from the client. We are setting the token as the value of Authorization header. If there is no such header, the token will be None. In that case, we set the value of user in g to None and return from the function. Refer to the previous chapters if you don’t remember the working of g.
We then try to decode the token to get the user id using the decode_token function we defined in the last chapter. If the decoding fails, we set the value of user in g to None and return from the function.
If the token decode is successful, we use the decoded user id to get the details from the database. After fetching the user details, we set it to g.user.

To restrict our APIs to all logged-in users, we can update each function to check for g.user and return an error if it is None. But we do not want to duplicate that across all our view functions. Fortunately, python offers us a more elegant solution. We can create a decorator and use it with all our view functions.

To start, add the following import to the top of auth.py file.

import functools

Then add the following code at the end of the file.

def login_required(view):
    @functools.wraps(view)
    def wrapped_view(**kwargs):
        if g.user is None:
            return {"error": "Authentication failed"}, 401

        return view(**kwargs)

    return wrapped_view

We created a login_required function that takes another function as its argument. This argument will be our view function.
functools.wraps decorator is used to preserve the attributes such as __name__ and __doc__ of the view function. Flask relies on these attributes.
Inside the wrapped function, we check if the user in g is None. If it is, we return an error with a status code 401.
If there is a value for user, we allow the normal request flow by calling the view function with all the arguments.

(Check out my tutorial on decorators if you want to know more about creating decorators.)

Next, we need to use this decorator on the APIs in the users blueprint. Open user.py and add an import at the top of the file for our new decorator.

from .auth import login_required

Then add the login_required decorator to our users_list function. This decorator should come after the blueprint decorator.

@blueprint.route("/", methods=["GET"])
@login_required
def list_users():
    # Rest of the code is not changed

Go ahead and add this decorator to get_user, update_user and delete_user functions. We don’t want to use this on create_user function since everyone should be able to create a new user account without any authentication.

To test this out, go to the headers section in postman and add a new header. Use Authorization as key and the token we got from the /auth/get-token/ as value. Hit send and you should see the regular API response. Try without the header and see what happens.

Additional Authorization for Update and Delete

We still have one more issue with our API permissions. Any user can edit or delete any other users' details. We don’t want that. Users should be able to update and delete their own profiles but they should not be able to do that in other profiles.

This is simple enough. We need to add the following condition to update_user and delete_user functions.

First, update the import statement in users.py to include g.

from flask import Blueprint, g, jsonify, request

Next, update the view function to check user ids.

@blueprint.route("/<int:user_id>/", methods=["PUT"])
@login_required
def update_user(user_id):
    if user_id != g.user["id"]:
        return {"error": "Permission denied"}, 401
    
    # Rest of the code unchanged

Here we are checking if the user id in the URL matches the id in the g.user. If it matches, a user is trying to update his details. Otherwise, we return an error. Update the delete_user function also with the same condition.

Congratulations! Now you know how to create REST APIs with python and flask.