JWT Authentication Service — Capstone Project for Token-Based Security

from flask import Flask, request, jsonify
import jwt, redis, datetime, secrets, hashlib

app = Flask(__name__)
r = redis.Redis(decode_responses=True)

# Configuration
HS256_SECRET = secrets.token_hex(32)
API_AUDIENCES = ["user-api", "threat-api", "device-api"]

# In-memory user store (use database in production)
USERS = {"admin": {"password": "pass123", "role": "admin"}}

# RS256 key pair (generate on startup)
# See jwt-signing lesson for key generation

# === YOUR IMPLEMENTATION ===

# 1. POST /api/v1/login
# 2. POST /api/v1/refresh
# 3. POST /api/v1/logout
# 4. POST /api/v1/oauth/token
# 5. GET /.well-known/jwks.json
# 6. GET /api/v1/me
# 7. Token validation middleware

def require_jwt(allowed_audiences=None):
    """Middleware that validates JWT and checks audience."""
    def decorator(f):
        def wrapper(*args, **kwargs):
            auth = request.headers.get("Authorization", "")
            if not auth.startswith("Bearer "):
                return jsonify({"error": "Missing token"}), 401

            token = auth[7:]
            # Determine algorithm from header
            header = jwt.get_unverified_header(token)
            alg = header.get("alg")

            try:
                if alg == "HS256":
                    payload = jwt.decode(
                        token, HS256_SECRET,
                        algorithms=["HS256"]
                    )
                elif alg == "RS256":
                    # Fetch key from JWKS
                    key = get_key_from_jwks(header.get("kid"))
                    if not key:
                        return jsonify({"error": "Unknown key"}), 401
                    payload = jwt.decode(
                        token, key, algorithms=["RS256"]
                    )
                else:
                    return jsonify({"error": "Unsupported algorithm"}), 401

                # Check blacklist
                jti = payload.get("jti")
                if jti and r.exists(f"blacklist:{jti}"):
                    return jsonify({"error": "Token revoked"}), 401

                # Check audience
                if allowed_audiences:
                    aud = payload.get("aud")
                    if aud not in allowed_audiences:
                        return jsonify({"error": "Wrong audience"}), 403

                g.user = payload

            except jwt.ExpiredSignatureError:
                return jsonify({"error": "Token expired"}), 401
            except jwt.InvalidTokenError:
                return jsonify({"error": "Invalid token"}), 401

            return f(*args, **kwargs)
        return wrapper
    return decorator

if __name__ == "__main__":
    app.run(port=5000)

# Test user login
curl -X POST -H "Content-Type: application/json" \
  -d '{"username":"admin","password":"pass123","audience":"user-api"}' \
  http://localhost:5000/api/v1/login

# Test service auth
curl -X POST -H "Content-Type: application/json" \
  -d '{"client_id":"threat-service","client_secret":"secret","audience":"threat-api"}' \
  http://localhost:5000/api/v1/oauth/token

# Fetch JWKS
curl http://localhost:5000/.well-known/jwks.json