Flask JWT Authentication

In this lesson, we'll explore how to implement JWT (JSON Web Token) authentication in Flask applications. JWT has become a standard for securing APIs and managing user sessions in modern web applications.

Introduction to JWT

JWT (JSON Web Token) is an open standard that defines a compact and self-contained way to securely transmit information between parties as a JSON object. This information can be verified and trusted because it is digitally signed.

Some key benefits of using JWT:

• Stateless: Servers don't need to store session information
• Portable: Same token can work across multiple backends
• Secure: Tokens are signed and can be encrypted
• Expirable: Tokens can have an expiration time

How JWT Works

A JWT token consists of three parts:

1. Header - Contains the algorithm used for signing
2. Payload - Contains claims (user data and metadata)
3. Signature - Verifies the token hasn't been tampered with

These parts are encoded and concatenated with dots to form a string like: xxxxx.yyyyy.zzzzz

Setting Up JWT Authentication in Flask

Let's implement JWT authentication in a Flask application step by step.

Step 1: Install Required Packages

First, we need to install the necessary packages:

pip install flask flask-jwt-extended

Step 2: Create a Basic Flask Application

Let's create a simple Flask application with JWT authentication:

from flask import Flask, jsonify, request
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity

# Create Flask app
app = Flask(__name__)

# Configure JWT
app.config['JWT_SECRET_KEY'] = 'super-secret-key'  # Change this in production!
jwt = JWTManager(app)

# Mock user database
users = {
    'user1': {
        'password': 'password1',
        'role': 'admin'
    },
    'user2': {
        'password': 'password2',
        'role': 'user'
    }
}

@app.route('/')
def index():
    return jsonify(message='Welcome to JWT Authentication API')

# Start the server
if __name__ == '__main__':
    app.run(debug=True)

Step 3: Implement Login Endpoint

Now let's create a login endpoint that verifies user credentials and returns a JWT token:

@app.route('/login', methods=['POST'])
def login():
    if not request.is_json:
        return jsonify({"error": "Missing JSON in request"}), 400

    username = request.json.get('username', None)
    password = request.json.get('password', None)

    if not username or not password:
        return jsonify({"error": "Missing username or password"}), 400

    # Check if user exists and password is correct
    if username not in users or users[username]['password'] != password:
        return jsonify({"error": "Invalid username or password"}), 401

    # Create access token with user identity
    access_token = create_access_token(identity=username)
    
    return jsonify(access_token=access_token), 200

Step 4: Create Protected Endpoints

Let's create endpoints that require JWT authentication:

@app.route('/protected', methods=['GET'])
@jwt_required()
def protected():
    # Access the identity of the current user
    current_user = get_jwt_identity()
    return jsonify(logged_in_as=current_user), 200

@app.route('/admin', methods=['GET'])
@jwt_required()
def admin():
    # Access the identity of the current user
    current_user = get_jwt_identity()
    
    # Check if user is admin
    if users[current_user]['role'] != 'admin':
        return jsonify({"error": "Admin access required"}), 403
        
    return jsonify(admin_access=True), 200

Step 5: Testing the Authentication Flow

Here's how to test the authentication flow using cURL:

1. Login to get a token:

curl -X POST -H "Content-Type: application/json" \
-d '{"username":"user1","password":"password1"}' \
http://localhost:5000/login

Output:

{
  "access_token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJmcmVzaCI6ZmFsc2UsImlhdCI6MTY0NzM2NjIwMiwianRpIjoiZjdkYzA2M2QtMjFkOC00MDZkLTk5YzktNDU5MmYxMTM4Y2I4IiwidHlwZSI6ImFjY2VzcyIsInN1YiI6InVzZXIxIiwibmJmIjoxNjQ3MzY2MjAyLCJleHAiOjE2NDczNjcxMDJ9.v41hh71VJpnnZBF9OaPvNVEmOI-qohbqhTdhCEPzmcs"
}

2. Access protected endpoint with token:

curl -H "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9..." \
http://localhost:5000/protected

Output:

{
  "logged_in_as": "user1"
}

Enhanced JWT Authentication Features

Let's explore some advanced features of Flask-JWT-Extended.

Token Expiration

You can configure token expiration time:

# Configure JWT with expiration time
app.config['JWT_SECRET_KEY'] = 'super-secret-key'
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = 3600  # 1 hour in seconds
jwt = JWTManager(app)

Refresh Tokens

Refresh tokens allow you to issue new access tokens without requiring the user to re-authenticate:

from flask_jwt_extended import create_refresh_token, jwt_refresh_token_required

@app.route('/login', methods=['POST'])
def login():
    # Authentication code...
    
    # Create both access and refresh 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

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

Token Revocation

You can implement token revocation to invalidate tokens before they expire:

from datetime import datetime
from flask_jwt_extended import get_jwt

# Store for revoked tokens
revoked_tokens = set()

@jwt.token_in_blocklist_loader
def check_if_token_revoked(jwt_header, jwt_payload):
    jti = jwt_payload["jti"]
    return jti in revoked_tokens

@app.route('/logout', methods=['DELETE'])
@jwt_required()
def logout():
    jti = get_jwt()["jti"]
    revoked_tokens.add(jti)
    return jsonify(message="Successfully logged out"), 200

Real-world Example: User Management API

Let's create a more comprehensive example with user registration, authentication, and protected resources:

from flask import Flask, jsonify, request
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
from werkzeug.security import generate_password_hash, check_password_hash
import datetime

app = Flask(__name__)

# Configure JWT
app.config['JWT_SECRET_KEY'] = 'change-this-key-in-production'
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = datetime.timedelta(hours=1)
jwt = JWTManager(app)

# In-memory user database (use a real database in production)
users_db = {}

# User registration endpoint
@app.route('/register', methods=['POST'])
def register():
    if not request.is_json:
        return jsonify({"error": "Missing JSON in request"}), 400

    username = request.json.get('username', None)
    password = request.json.get('password', None)
    email = request.json.get('email', None)

    if not username or not password or not email:
        return jsonify({"error": "Missing required fields"}), 400
    
    if username in users_db:
        return jsonify({"error": "Username already exists"}), 409
    
    # Hash the password for security
    hashed_password = generate_password_hash(password)
    
    # Store new user
    users_db[username] = {
        'password': hashed_password,
        'email': email,
        'created_at': datetime.datetime.now().isoformat()
    }
    
    return jsonify({"message": "User created successfully"}), 201

# User login endpoint
@app.route('/login', methods=['POST'])
def login():
    if not request.is_json:
        return jsonify({"error": "Missing JSON in request"}), 400

    username = request.json.get('username', None)
    password = request.json.get('password', None)

    if not username or not password:
        return jsonify({"error": "Missing username or password"}), 400

    # Check if user exists
    if username not in users_db:
        return jsonify({"error": "Invalid username or password"}), 401
    
    # Check if password is correct
    if not check_password_hash(users_db[username]['password'], password):
        return jsonify({"error": "Invalid username or password"}), 401

    # Create access token with user identity
    access_token = create_access_token(identity=username)
    
    return jsonify(access_token=access_token), 200

# User profile endpoint (protected)
@app.route('/profile', methods=['GET'])
@jwt_required()
def profile():
    current_user = get_jwt_identity()
    
    if current_user not in users_db:
        return jsonify({"error": "User not found"}), 404
    
    # Don't return password hash
    user_data = {
        'username': current_user,
        'email': users_db[current_user]['email'],
        'created_at': users_db[current_user]['created_at']
    }
    
    return jsonify(user=user_data), 200

if __name__ == '__main__':
    app.run(debug=True)

Best Practices for JWT Authentication

When implementing JWT authentication in production applications, follow these best practices:

1. Use HTTPS: Always serve your API over HTTPS to prevent token theft via network sniffing.

2. Secure Secret Keys: Store your JWT secret keys securely, preferably using environment variables.

import os
app.config['JWT_SECRET_KEY'] = os.environ.get('JWT_SECRET_KEY', 'default-dev-key')

3. Short Token Expiration: Use short expiration times for access tokens (minutes to hours).

4. Token Refresh Strategy: Implement refresh tokens to improve user experience without compromising security.

5. Include Only Necessary Claims: Don't store sensitive information in the token payload.

6. Token Revocation: Implement token revocation for logout functionality.

7. Handle Errors Gracefully: Provide clear error messages without exposing sensitive details.

@jwt.expired_token_loader
def expired_token_callback(jwt_header, jwt_payload):
    return jsonify({
        'status': 401,
        'sub_status': 42,
        'msg': 'The token has expired'
    }), 401

Common JWT Issues and Solutions

Issue: CORS with JWT

When building SPAs with JWT authentication, you might encounter CORS issues:

from flask_cors import CORS

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})  # In production, specify origins

Issue: JWT Token Storage on Client Side

Client-side storage recommendations:

• Store in HttpOnly cookies for web apps
• For SPAs, use memory storage with refresh token in HttpOnly cookie

Issue: Handling Expired Tokens

Implement client-side logic to handle token expiration:

• Detect 401 responses
• Use refresh token to get a new access token
• Redirect to login if refresh fails

Summary

In this lesson, we've covered:

1. Introduction to JWT and its components
2. Setting up JWT authentication in Flask using Flask-JWT-Extended
3. Creating login and protected endpoints
4. Advanced features like token expiration, refresh tokens, and revocation
5. Building a real-world user management API
6. Best practices for secure JWT implementation

JWT authentication provides a secure, stateless way to authenticate users in your Flask applications, making it ideal for modern web APIs and single-page applications.

Additional Resources

• Flask-JWT-Extended Documentation
• JWT.io - Decode, verify and generate JWT tokens
• OWASP JWT Cheat Sheet

Exercises

1. Modify the example to use a database like SQLAlchemy instead of in-memory storage.
2. Add role-based access control to your API endpoints.
3. Implement token blacklisting using Redis for better scalability.
4. Create a frontend that interacts with the JWT authentication API.
5. Add JWT claims customization to include user permissions in the token.