DEV Community

Tooleroid
Tooleroid

Posted on

JWT Authentication in Python: FastAPI and Flask Implementation Guide

Python's rich ecosystem provides multiple robust frameworks for implementing JWT authentication. In this comprehensive guide, we'll explore how to implement secure JWT authentication using two of Python's most popular web frameworks: FastAPI and Flask. We'll cover everything from basic setup to advanced security features, with a focus on real-world applications and best practices.

Why Choose Python for JWT Authentication?

Python offers several advantages for implementing JWT authentication:

  1. Rich Framework Ecosystem: Frameworks like FastAPI and Flask provide built-in support for JWT authentication.
  2. Strong Security Libraries: Python's security libraries are well-maintained and regularly updated.
  3. Type Hints Support: Modern Python's type hints help catch authentication-related bugs early.
  4. Excellent Documentation: Both FastAPI and Flask have comprehensive documentation for security features.

FastAPI Implementation

FastAPI is a modern, fast web framework that's perfect for building authenticated APIs. It provides built-in support for OAuth2 with JWT tokens and automatic API documentation.

Basic Setup

First, let's set up our FastAPI application with the necessary dependencies and models:

from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
from pydantic import BaseModel

# Models
class Token(BaseModel):
    access_token: str
    token_type: str
    expires_in: int

class TokenData(BaseModel):
    sub: Optional[str] = None
    exp: Optional[int] = None

class User(BaseModel):
    username: str
    email: str
    disabled: Optional[bool] = None

# Configuration
SECRET_KEY = "your-secret-key"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

app = FastAPI()
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
Enter fullscreen mode Exit fullscreen mode

This setup provides several important features:

  • Type-safe models using Pydantic
  • Built-in OAuth2 password flow
  • Automatic OpenAPI documentation
  • Configurable token expiration

Token Generation

def create_access_token(data: Dict[str, Any], expires_delta: timedelta) -> str:
    to_encode = data.copy()
    expire = datetime.utcnow() + expires_delta
    to_encode.update({
        "exp": expire,
        "iat": datetime.utcnow()
    })
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

async def get_current_user(token: str = Depends(oauth2_scheme)) -> User:
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(sub=username)
    except JWTError:
        raise credentials_exception

    user = get_user(username=token_data.sub)
    if user is None:
        raise credentials_exception
    return user
Enter fullscreen mode Exit fullscreen mode

Key security considerations in token handling:

  • Use of cryptographically secure algorithms
  • Proper error handling for invalid tokens
  • Automatic token expiration
  • Type-safe token validation

FastAPI Routes

@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    user = authenticate_user(form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username},
        expires_delta=access_token_expires
    )

    return {
        "access_token": access_token,
        "token_type": "bearer",
        "expires_in": ACCESS_TOKEN_EXPIRE_MINUTES * 60
    }

@app.get("/users/me", response_model=User)
async def read_users_me(current_user: User = Depends(get_current_user)):
    return current_user
Enter fullscreen mode Exit fullscreen mode

Flask Implementation

Basic Setup

from flask import Flask, jsonify, request
from flask_jwt_extended import (
    JWTManager, create_access_token,
    jwt_required, get_jwt_identity
)
from datetime import timedelta

app = Flask(__name__)

# Setup Flask-JWT-Extended
app.config["JWT_SECRET_KEY"] = "your-secret-key"
app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(hours=1)
jwt = JWTManager(app)
Enter fullscreen mode Exit fullscreen mode

Token Management

class TokenManager:
    def __init__(self):
        self.blacklist = set()

    def add_to_blacklist(self, jti):
        self.blacklist.add(jti)

    def is_blacklisted(self, jti):
        return jti in self.blacklist

token_manager = TokenManager()

@jwt.token_in_blocklist_loader
def check_if_token_revoked(jwt_header, jwt_payload):
    jti = jwt_payload["jti"]
    return token_manager.is_blacklisted(jti)
Enter fullscreen mode Exit fullscreen mode

Flask Routes

@app.route("/login", methods=["POST"])
def login():
    username = request.json.get("username", None)
    password = request.json.get("password", None)

    user = authenticate_user(username, password)
    if not user:
        return jsonify({"error": "Invalid credentials"}), 401

    access_token = create_access_token(
        identity=username,
        additional_claims={
            "roles": user.roles,
            "email": user.email
        }
    )

    return jsonify({
        "access_token": access_token,
        "token_type": "bearer",
        "expires_in": 3600
    })

@app.route("/protected", methods=["GET"])
@jwt_required()
def protected():
    current_user = get_jwt_identity()
    return jsonify(logged_in_as=current_user), 200
Enter fullscreen mode Exit fullscreen mode

Security Best Practices

Rate Limiting

from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

limiter = Limiter(
    app=app,
    key_func=get_remote_address,
    default_limits=["200 per day", "50 per hour"]
)

@app.route("/login", methods=["POST"])
@limiter.limit("5 per minute")
def login():
    # ... login logic
    pass
Enter fullscreen mode Exit fullscreen mode

Refresh Token Implementation

class RefreshTokenStore:
    def __init__(self):
        self.refresh_tokens = {}

    def store_refresh_token(self, user_id: str, refresh_token: str):
        self.refresh_tokens[user_id] = refresh_token

    def validate_refresh_token(self, user_id: str, refresh_token: str) -> bool:
        stored_token = self.refresh_tokens.get(user_id)
        return stored_token == refresh_token

    def revoke_refresh_token(self, user_id: str):
        self.refresh_tokens.pop(user_id, None)

@app.route("/refresh", methods=["POST"])
@jwt_required(refresh=True)
def refresh():
    identity = get_jwt_identity()
    access_token = create_access_token(identity=identity)
    return jsonify(access_token=access_token)
Enter fullscreen mode Exit fullscreen mode

Error Handling

class AuthError(Exception):
    def __init__(self, error: str, status_code: int):
        super().__init__()
        self.error = error
        self.status_code = status_code

@app.errorhandler(AuthError)
def handle_auth_error(error):
    return jsonify({
        "error": error.error,
        "message": "Authentication failed"
    }), error.status_code

@jwt.expired_token_loader
def expired_token_callback(jwt_header, jwt_payload):
    return jsonify({
        "error": "token_expired",
        "message": "The token has expired"
    }), 401
Enter fullscreen mode Exit fullscreen mode

Testing

import pytest
from fastapi.testclient import TestClient

def test_login_success(client: TestClient):
    response = client.post(
        "/token",
        data={
            "username": "testuser",
            "password": "testpass"
        }
    )
    assert response.status_code == 200
    assert "access_token" in response.json()

def test_protected_route(client: TestClient, valid_token):
    response = client.get(
        "/users/me",
        headers={"Authorization": f"Bearer {valid_token}"}
    )
    assert response.status_code == 200
Enter fullscreen mode Exit fullscreen mode

Related Resources

Conclusion

Python's FastAPI and Flask frameworks provide robust solutions for implementing JWT authentication. By following these patterns and security practices, you can build secure and scalable authentication systems.

Remember to check our other security guides and authentication tools for more resources!

This guide provides comprehensive coverage of JWT implementation in Python, focusing on both FastAPI and Flask frameworks.

Use 400+ completely free and online tools at Tooleroid.com!

Top comments (0)