Flask and Keycloak for Single Sign On

Posted on Feb 14, 2025

Flask and Keycloak for JWT based Authentication

To secure your Flask API endpoints using Keycloak and JWT tokens, you can create a custom decorator that validates the JWT token before allowing access to the endpoint. Below is a step-by-step guide to achieve this:

1. Install Required Libraries

You need the following Python libraries:

Flask: For building the API.
PyJWT: For decoding and validating JWT tokens.
requests: For fetching the public key from Keycloak to validate the token.

Install them using pip:

pip install Flask PyJWT requests

2. Set Up Keycloak

Create a realm, client, and roles in Keycloak.
Ensure your Keycloak server is running and accessible.
Note down the following details:
- Keycloak server URL (e.g., http://localhost:8080/auth).
- Realm name.
- Client ID.
- Client secret (if using confidential clients).

3. Fetch Keycloak’s Public Key

Keycloak signs JWT tokens using a private key, and you can validate them using the corresponding public key. Fetch the public key from Keycloak’s certs endpoint:

GET http://<keycloak-url>/auth/realms/<realm-name>/protocol/openid-connect/certs

You can fetch this programmatically in your Flask app.

4. Create the Decorator

Here’s how you can create a decorator to validate the JWT token:

from functools import wraps
import jwt
from flask import request, jsonify
import requests

# Keycloak configuration
KEYCLOAK_URL = "http://localhost:8080/auth"
REALM_NAME = "your-realm"
CLIENT_ID = "your-client-id"
PUBLIC_KEY_URL = f"{KEYCLOAK_URL}/realms/{REALM_NAME}/protocol/openid-connect/certs"

# Fetch the public key from Keycloak
def get_public_key():
    response = requests.get(PUBLIC_KEY_URL)
    response.raise_for_status()
    jwks = response.json()
    public_key = jwt.algorithms.RSAAlgorithm.from_jwk(jwks['keys'][0])
    return public_key

# Decorator to validate JWT token
def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None

        # Get the token from the Authorization header
        if 'Authorization' in request.headers:
            token = request.headers['Authorization'].split(" ")[1]

        if not token:
            return jsonify({"message": "Token is missing!"}), 401

        try:
            # Decode and validate the token
            public_key = get_public_key()
            data = jwt.decode(token, public_key, algorithms=["RS256"], audience=CLIENT_ID)
            # Optionally, you can add more checks here (e.g., roles, expiration, etc.)
        except jwt.ExpiredSignatureError:
            return jsonify({"message": "Token has expired!"}), 401
        except jwt.InvalidTokenError:
            return jsonify({"message": "Invalid token!"}), 401

        # Attach the decoded token data to the request object
        request.user_data = data
        return f(*args, **kwargs)

    return decorated

5. Use the Decorator in Your Flask App

Apply the @token_required decorator to your API endpoints:

from flask import Flask, jsonify

app = Flask(__name__)

@app.route('/protected', methods=['GET'])
@token_required
def protected():
    return jsonify({"message": "This is a protected endpoint!", "user": request.user_data})

@app.route('/unprotected', methods=['GET'])
def unprotected():
    return jsonify({"message": "This is an unprotected endpoint!"})

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

6. Disable Flask’s Default Session Management

Flask uses browser sessions by default. To disable this and rely solely on JWT tokens, ensure you don’t use flask.session or any session-related functionality in your app.

7. Test Your API

Use a tool like Postman or curl to test your API.
Obtain a JWT token from Keycloak (e.g., via the /protocol/openid-connect/token endpoint).
Include the token in the Authorization header of your requests:
```
Authorization: Bearer <your-jwt-token>
```

8. Optional: Add Role-Based Access Control

If you want to enforce role-based access, you can extend the decorator to check for specific roles in the token’s realm_access or resource_access claims.

Example:

def role_required(role):
    def decorator(f):
        @wraps(f)
        def decorated(*args, **kwargs):
            if 'realm_access' not in request.user_data or role not in request.user_data['realm_access']['roles']:
                return jsonify({"message": "Access denied! Role required."}), 403
            return f(*args, **kwargs)
        return decorated
    return decorator

@app.route('/admin', methods=['GET'])
@token_required
@role_required('admin')
def admin():
    return jsonify({"message": "Welcome, admin!"})

This setup ensures your Flask API is secured using Keycloak and JWT tokens, without relying on Flask’s default session management.