Using JWTs in Python Flask REST Framework

JSON Web Tokens (JWTs) secure communication between parties over the internet by authenticating users and transmitting information securely, without requiring a centralized storage system.

In this article, we'll explain what JWTs are and give a high-level overview of how they work. We'll also implement a JWT-based authentication system by creating a to-do list API using Flask.

So, whether you're a beginner or an experienced Python developer, this guide aims to enhance your understanding of JWTs and their practical application in Flask.

First, let's provide an overview of JWTs and explain how they work, at a high level.

JWTs: How They Work

JWTs are a compact and self-contained method for securely transmitting information between parties as a JSON object. This information can be verified and trusted because it is digitally signed using a secret or a public/private key pair.

A JWT is composed of three parts:

• Header: Contains metadata about the token, including the type of token and the hashing algorithm used to sign it. Typical cryptographic algorithms include HMAC with SHA-256 (HS256) and an RSA signature with SHA-256 (RS256). Here's an example of how this might look:

  JSON

  {
    "alg": "HS256",
    "typ": "JWT"
  }

  In this case, alg specifies the hashing algorithm used, while typ indicates the token type (which, of course, is "JWT"). Read about other commonly used header fields.

• Payload: Contains the claims, which are statements about an entity (typically, the user) and additional data that can include information like the user ID, username, and roles. Here's an example of how the payload might look:

  JSON

  {
    "sub": "1234567890",
    "name": "John Doe",
    "admin": true,
    "iat": 1516239022
  }

  In this case, sub is the subject identifier and represents the subject of the token (e.g., the user ID); name is the name of the user; admin is a boolean indicating administrative privileges; iat means "issued at time" and represents when the token was issued (in the Unix timestamp).

• Signature: The signature is used to verify the authenticity of the token and to ensure that it hasn't been tampered with. It is created by encoding the header and payload, concatenating them with a period, and then hashing them using the algorithm specified in the header along with a secret key.

So, when putting everything together, the structure of a JWT looks like this:

{Header in Base64URL}.{Payload in Base64URL}.{Signature}

When decoded, the JWT reveals the original JSON structures:

{
  "header": {
    "alg": "HS256",
    "typ": "JWT"
  },
  "payload": {
    "sub": "1234567890",
    "name": "John Doe",
    "admin": true,
    "iat": 1516239022
  },
  "signature": "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
}

Why JWTs Are Useful for Secure User Authentication in APIs

JWTs offer several benefits when securing user authentication in APIs, including:

• Stateless sessions: JWTs eliminate the need to store session information on the server, as all the necessary data is contained within a token. This makes scaling your application easier.
• Security: Because JWTs are signed, the server can verify the integrity and authenticity of the token. This ensures that the data hasn't been altered and comes from a trusted source. This also helps protect endpoints by ensuring that only authenticated users can access certain resources.
• Efficiency: JWTs are compact, making them ideal for inclusion in HTTP headers and efficient for network transmission.
• Flexibility: They can include a variety of claims, allowing you to convey user roles, permissions, and other metadata necessary for authorization.
• Cross-domain support: JWTs can be used across different domains without issues, making them suitable for Single Sign-On (SSO) systems.
• Simplified authentication flow: They streamline the authentication process, reducing the complexity of managing user sessions.

Setting Up the Flask Environment for a Flask Python App

Let's now dive into a practical example to show how to implement theory into practice.

But before that, let's first set up a Flask environment with all the necessary dependencies.

Installing Dependencies and Project Structure

First, make sure you have Python 3.6+ installed. Then, create a project directory like this one:

jwt_flask_todo/
├── app.py
└── venv/

Where:

• app.py is the main Flask application file.
• venv/ is the virtual environment where dependencies will be installed.

To create the virtual environment, navigate to the root folder of your project and execute the command below:

python3 -m venv venv

Then activate it:

Linux/MacOS:

source ./venv/bin/activate

Windows:

.\venv\Scripts\activate.bat

You can now install the libraries needed to replicate the Flask app:

pip install Flask Flask-JWT-Extended

Note that Flask-JWT-Extended is the library used in the following code examples to manage JWT tokens.

Implementing JWT Authentication in a Flask API

We'll start our to-do list API by implementing user registration and login functionality, integrating JWTs for authentication.

Building a User Registration System

Let's first create an endpoint to register users, storing their data in an in-memory dictionary for simplicity:

from flask import Flask, request, jsonify
 
app = Flask(__name__)
 
# In-memory 'database' of users
users = {}
 
@app.route('/register', methods=['POST'])
def register():
    username = request.json.get('username')
    password = request.json.get('password')
 
    # Input validation
    if not username or not password:
        return jsonify({"msg": "Missing username or password"}), 400
 
    # Check if user already exists
    if username in users:
        return jsonify({"msg": "User already exists"}), 400
 
    # Store user
    users[username] = password
    return jsonify({"msg": "User registered successfully"}), 201
 
if __name__ == '__main__':
    app.run(debug=True)

This endpoint does the following:

• Data retrieval: Extracts a username and password from the JSON request body.
• Validation: Checks if the username and password are provided.
• User existence check: Ensures the username isn't already taken.
• User storage: Adds the new user to the users dictionary.
• Response: Returns a success message if everything works as expected.

Adding User Login with JWTs

Next, let's create a login endpoint that issues a JWT token when a user logs in successfully:

from flask_jwt_extended import JWTManager, create_access_token
 
# Configure the JWT secret key
app.config['JWT_SECRET_KEY'] = 'your_jwt_secret_key'
jwt = JWTManager(app)
 
@app.route('/login', methods=['POST'])
def login():
    username = request.json.get('username')
    password = request.json.get('password')
 
    # Input validation
    if not username or not password:
        return jsonify({"msg": "Missing username or password"}), 400
 
    # Authentication
    if users.get(username) != password:
        return jsonify({"msg": "Invalid credentials"}), 401
 
    # Create JWT
    access_token = create_access_token(identity=username)
    return jsonify(access_token=access_token), 200

This endpoint is similar to the user registration one. In addition, it does the following:

• Authentication: Checks if the provided password matches the stored one.
• Token creation: Generates an access token with the create_access_token method, embedding the username as the identity.

Protecting To-Do List API Endpoints with JWTs

Now that you have implemented a user registration and login, you can secure the endpoints so that only authenticated users can access them.

Apply the @jwt_required() decorator to the routes you want to protect. Use get_jwt_identity() to retrieve the current user's identity from the token:

from flask_jwt_extended import JWTManager, jwt_required, get_jwt_identity
 
# Initialize the JWTManager:
app = Flask(__name__)
app.config['JWT_SECRET_KEY'] = 'your_jwt_secret_key'
jwt = JWTManager(app)
 
# Secure the endpoint
@app.route('/protected', methods=['GET'])
@jwt_required()
def protected():
    # Access the identity of the current user w
    current_user = get_jwt_identity()
    return jsonify(logged_in_as=current_user), 200

So, here's what this code does:

• @jwt_required() decorator: Ensures that the endpoint can only be accessed if a valid JWT is present in the request. If the JWT is missing or invalid, the request will be rejected.
• get_jwt_identity() function: Inside the protected route, this function retrieves a current user's identity from the JWT. The identity is the value you set when creating the token (usually, the username or user ID).

Creating a To-Do List API

Let's now build a simple to-do list API where each user can manage their tasks. We'll create endpoints to manage the classical CRUD operations we'd expect to see in such an app.

Creating a Task

First, let's create a task:

@app.route('/tasks', methods=['POST'])
@jwt_required()
def add_task():
    current_user = get_jwt_identity()
    task = request.json.get('task')
 
    if not task:
        return jsonify({"msg": "Task content missing"}), 400
 
    # Initialize user's task list if it doesn't exist
    if current_user not in to_do_lists:
        to_do_lists[current_user] = []
 
    # Add task to user's task list
    to_do_lists[current_user].append(task)
    return jsonify({"msg": "Task added", "task": task}), 201

This endpoint does the following:

• @jwt_required(): Ensures that only authenticated users can access the endpoint.
• get_jwt_identity(): Retrieves the current user's username from the JWT.
• Task retrieval: Gets the task content from the request JSON.
• Validation: Checks if the task content is provided.
• Task list initialization: Creates an empty list for the user if they don't have one.
• Adding task: Appends the task to the user's list.
• Response: Returns a confirmation message with the task added.

Viewing All Tasks

This endpoint fetches the task list for a current user and returns an empty list if none exists:

@app.route('/tasks', methods=['GET'])
@jwt_required()
def get_tasks():
    current_user = get_jwt_identity()
    tasks = to_do_lists.get(current_user, [])
    return jsonify({"tasks": tasks}), 200

Updating a Task

Now we'll update a task:

@app.route('/tasks/<int:task_id>', methods=['PUT'])
@jwt_required()
def update_task(task_id):
    current_user = get_jwt_identity()
    tasks = to_do_lists.get(current_user, [])
 
    if task_id < 0 or task_id >= len(tasks):
        return jsonify({"msg": "Task not found"}), 404
 
    updated_task = request.json.get('task')
    if not updated_task:
        return jsonify({"msg": "Task content missing"}), 400
 
    tasks[task_id] = updated_task
    return jsonify({"msg": "Task updated", "task": updated_task}), 200

In this endpoint:

• task_id parameter: Represents the index of the task in the user's task list.
• Index validation: Ensures the task ID is within valid bounds.
• Updated task content: Retrieves the new task content from the request.
• Updating task: Replaces the existing task with the updated content.
• Response: Confirms the update with the new task.

Deleting a Task

Finally, this endpoint removes the task at the specified index:

@app.route('/tasks/<int:task_id>', methods=['DELETE'])
@jwt_required()
def delete_task(task_id):
    current_user = get_jwt_identity()
    tasks = to_do_lists.get(current_user, [])
 
    if task_id < 0 or task_id >= len(tasks):
        return jsonify({"msg": "Task not found"}), 404
 
    deleted_task = tasks.pop(task_id)
    return jsonify({"msg": "Task deleted", "task": deleted_task}), 200

For longer-lived sessions, you can allow users to refresh their tokens instead of having to log in again after every expiration.

Implementing Token Refresh for Longer Sessions

You can set up token refreshes for longer sessions in a few different ways:

• Update the JWT configuration.
• Modify the login endpoint to issue refresh tokens.
• Create a new endpoint that manages token refreshes.

Let's take each approach in turn.

Updating JWT Configuration

JWTs expiration time can be managed like so:

from datetime import timedelta
 
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(minutes=15)  # Access tokens expire after 15 minutes
app.config['JWT_REFRESH_TOKEN_EXPIRES'] = timedelta(days=30)    # Refresh tokens expire after 30 days

Modifying the Login Endpoint to Issue Refresh Tokens

To use both access and refresh tokens, you can update the login endpoint as follows:

from flask_jwt_extended import create_refresh_token
 
@app.route('/login', methods=['POST'])
def login():
    # ... [existing code]
 
    # Authentication
    if users.get(username) != password:
        return jsonify({"msg": "Invalid credentials"}), 401
 
    # Create tokens
    access_token = create_access_token(identity=username)
    refresh_token = create_refresh_token(identity=username)
    return jsonify(access_token=access_token, refresh_token=refresh_token), 200

In this case, you can use the create_refresh_token() method to generate a refresh token for the user and prolong their session.

Adding Token Refresh Endpoint

As in the last case, you can create an endpoint to refresh access tokens using the refresh token:

from flask_jwt_extended import jwt_refresh_token_required
 
@app.route('/refresh', methods=['POST'])
@jwt_refresh_token_required
def refresh():
    current_user = get_jwt_identity()
    new_access_token = create_access_token(identity=current_user)
    return jsonify(access_token=new_access_token), 200

Here, though, the @jwt_refresh_token_required decorator ensures the endpoint is accessed with a valid refresh token. Then, the create_access_token() method creates a new access token to refresh the endpoint.

Logging Out and Managing User Sessions

Since JWTs are stateless, they cannot be invalidated server-side once issued. For this reason, you can implement token blocklisting to simulate a logout.

Let's see how.

Blocklisting Tokens

Revoked tokens can be maintained into a blocklist like so:

from flask_jwt_extended import get_raw_jwt
 
# Set to store blocklisted tokens
blocklisted_tokens = set()
 
@app.route('/logout', methods=['DELETE'])
@jwt_required()
def logout():
    jti = get_raw_jwt()['jti']
    blocklisted_tokens.add(jti)
    return jsonify({"msg": "Successfully logged out"}), 200

The get_raw_jwt() method retrieves the JWT data, including the jti (which is the JWT ID). Then, this endpoint adds the token's jti to the blocklisted_tokens set, simulating a logout. Once a token is inserted into the blocklist, any attempt to use it to log in will fail.

So, before processing any request, you could check if the token a user is trying to use has been revoked. This could be done like so:

from flask_jwt_extended import jwt_required, get_jwt_claims
 
@app.before_request
def check_revoked_token():
    jti = get_raw_jwt().get('jti')
    if jti in blocklisted_tokens:
        return jsonify({"msg": "Token revoked"}), 401

So, this piece of code can run before every request and checks if the token's jti is in the blocklist. If it is, it returns a response indicating that the token has been revoked.

Adding Permissions and Differentiating Between User Roles

To enhance security, you can manage control access to certain endpoints. This can be achieved, for example, through Role-Based Access Control (RBAC).

So, let's see how you can implement a simple RBAC system that differentiates between standard users and admins. But before that, bear in mind that in this article:

• Standard users can only manage their own tasks.
• Admin users have full access, including the ability to delete any task.

Modifying User Registration to Include Roles

First, you have to modify a registration endpoint to allow setting a user role. This can be done like so:

@app.route('/register', methods=['POST'])
def register():
    # Extract data from the request
    username = request.json.get('username')
    password = request.json.get('password')
    role = request.json.get('role', 'user')  # Default role is 'user'
 
    # Input validation
    if not username or not password:
        return jsonify({"msg": "Missing username or password"}), 400
 
    # Check if user already exists
    if username in users:
        return jsonify({"msg": "User already exists"}), 400
 
    # Store user with role
    users[username] = {'password': password, 'role': role}
    return jsonify({"msg": "User registered successfully"}), 201

Here's what's happening in this endpoint now:

• Role assignment: Now there's the possibility of setting a role during registration. In particular, note that the default assigned value is user if no role is explicitly provided.
• User data structure: The code stores user information as a dictionary containing the password and role. In other words, it stores the user with its own role.

Protecting Endpoints Based on Roles

To improve endpoint protection even further, you can protect based on users' roles.

Create a custom decorator that checks a user's role before allowing access to certain endpoints, like so:

from functools import wraps
from flask import abort
 
def role_required(required_role):
    def decorator(f):
        @wraps(f)
        @jwt_required()
        def wrapper(*args, **kwargs):
            claims = get_jwt()
            user_role = claims.get('role', 'user')
            if user_role != required_role:
                abort(403, description="Forbidden: Insufficient privileges")
            return f(*args, **kwargs)
        return wrapper
    return decorator

In this example, the role_required decorator checks if the user's role matches the required_role (the one needed to get access). If the user's role doesn't match, the endpoint aborts with a 403 Forbidden error message, meaning that the requester doesn't have the right privileges to authenticate the endpoint.

Finally, we'll take a brief look at setting permissions so that only admins can delete tasks.

Specific Operations: Only Admins Can Delete Tasks

There are cases where deleting something should only be allowed for administrators.

This can be implemented in the to-do list app we created by modifying the delete task endpoint, allowing only admins to delete tasks. Here's how:

from flask_jwt_extended import get_jwt
 
@app.route('/tasks/<int:task_id>', methods=['DELETE'])
@role_required('admin')
def delete_task(task_id):
    # Get all tasks (for all users)
    all_tasks = [task for tasks in to_do_lists.values() for task in tasks]
 
    # Check if task exists
    if task_id < 0 or task_id >= len(all_tasks):
        return jsonify({"msg": "Task not found"}), 404
 
    # Find and delete the task
    for user_tasks in to_do_lists.values():
        if task_id < len(user_tasks):
            deleted_task = user_tasks.pop(task_id)
            return jsonify({"msg": "Task deleted by admin", "task": deleted_task}), 200
        else:
            task_id -= len(user_tasks)

So now, the delete task endpoint uses the @role_required('admin') decorator to ensure only admins can access this endpoint. This means that only admins are allowed to delete any task from any user. Then, it confirms that the deletion is done.

And that's it!

Wrapping Up

By creating a to-do list API, we demonstrated how JWTs can secure user data and ensure that only authenticated users can access or modify their tasks.

We covered registration, login, task management, token refresh, and logout functionality, giving you the foundation to expand these concepts further!

Happy coding!