JWT & API Authentication

import secrets
from datetime import timedelta

from flask import Flask, jsonify, request
from werkzeug.security import check_password_hash, generate_password_hash
from flask_jwt_extended import (
    JWTManager,
    create_access_token,
    create_refresh_token,
    current_user,
    get_jwt,
    get_jwt_identity,
    jwt_required,
)

# A throwaway password for the demo; real ones come from the registration form.
DEMO_PASSWORD = secrets.token_urlsafe(12)
USERS = {
    "1": {
        "id": "1",
        "email": "ana@example.com",
        "password_hash": generate_password_hash(DEMO_PASSWORD),
        "roles": ["reader"],
    },
}
# Compared when the e-mail is unknown, so a missing user costs as much time as a wrong password.
DUMMY_HASH = generate_password_hash(secrets.token_urlsafe(12))
REVOKED_JTIS: set[str] = set()   # stands in for Redis or a database table, see the revocation section


def create_jwt_app() -> Flask:
    app = Flask(__name__)
    app.config["JWT_SECRET_KEY"] = secrets.token_hex(32)          # in production: read from the environment
    app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(minutes=15)
    app.config["JWT_REFRESH_TOKEN_EXPIRES"] = timedelta(days=7)
    jwt = JWTManager(app)

    @jwt.user_lookup_loader
    def load_user(jwt_header, jwt_data):
        return USERS.get(jwt_data["sub"])             # None makes current_user fail with a 401

    @jwt.token_in_blocklist_loader
    def is_revoked(jwt_header, jwt_data):
        return jwt_data["jti"] in REVOKED_JTIS

    @app.post("/api/login")
    def login():
        data = request.get_json(silent=True) or {}
        user = next((u for u in USERS.values() if u["email"] == data.get("email")), None)
        stored = user["password_hash"] if user else DUMMY_HASH
        if not check_password_hash(stored, data.get("password", "")) or user is None:
            return jsonify(error="bad credentials"), 401
        return jsonify(
            access_token=create_access_token(
                identity=user["id"], fresh=True, additional_claims={"roles": user["roles"]}
            ),
            refresh_token=create_refresh_token(identity=user["id"]),
        )

    @app.get("/api/me")
    @jwt_required()
    def me():
        return jsonify(id=get_jwt_identity(), email=current_user["email"], roles=get_jwt()["roles"])

    return app


app = create_jwt_app()
client = app.test_client()
tokens = client.post("/api/login", json={"email": "ana@example.com", "password": DEMO_PASSWORD}).get_json()
print(sorted(tokens))
print(client.get("/api/me").status_code, client.get("/api/me").get_json())
bearer = {"Authorization": f"Bearer {tokens['access_token']}"}
print(client.get("/api/me", headers=bearer).get_json())
print(client.post("/api/login", json={"email": "ana@example.com", "password": secrets.token_urlsafe(8)}).status_code)

curl -s -X POST http://localhost:5000/api/login \
  -H 'Content-Type: application/json' \
  -d '{"email": "ana@example.com", "password": "<password>"}'

curl -s http://localhost:5000/api/me -H 'Authorization: Bearer <access-token>'

{"email": "ana@example.com", "id": "1", "roles": ["reader"]}

from flask import make_response


def create_json_error_app() -> Flask:
    app = create_jwt_app()
    jwt = app.extensions["flask-jwt-extended"]

    def failure(code: str, status: int = 401):
        body = make_response(jsonify(error=code), status)
        body.headers["WWW-Authenticate"] = f'Bearer error="{code}"'
        return body

    @jwt.unauthorized_loader
    def missing_token(reason):                  # no Authorization header at all
        response = jsonify(error="authentication_required")
        response.status_code = 401
        response.headers["WWW-Authenticate"] = "Bearer"
        return response

    @jwt.invalid_token_loader
    def invalid_token(reason):                  # malformed, wrong signature, wrong type
        return failure("invalid_token")

    @jwt.expired_token_loader
    def expired_token(jwt_header, jwt_data):
        return failure("token_expired")

    @jwt.revoked_token_loader
    def revoked_token(jwt_header, jwt_data):
        return failure("token_revoked")

    return app


app = create_json_error_app()
client = app.test_client()
with app.app_context():
    expired = create_access_token(identity="1", expires_delta=timedelta(seconds=-1), additional_claims={"roles": []})
print(client.get("/api/me").get_json())
print(client.get("/api/me", headers={"Authorization": "Bearer not-a-token"}).status_code)
response = client.get("/api/me", headers={"Authorization": f"Bearer {expired}"})
print(response.status_code, response.headers["WWW-Authenticate"])

def create_refresh_app() -> Flask:
    app = create_json_error_app()

    @app.post("/api/refresh")
    @jwt_required(refresh=True)
    def refresh():
        claims = get_jwt()
        REVOKED_JTIS.add(claims["jti"])                                  # a refresh token is used once
        identity = get_jwt_identity()
        user = USERS[identity]
        return jsonify(
            access_token=create_access_token(identity=identity, additional_claims={"roles": user["roles"]}),
            refresh_token=create_refresh_token(identity=identity),
        )

    @app.post("/api/password")
    @jwt_required(fresh=True)
    def change_password():
        return jsonify(changed=True)

    @app.post("/api/logout")
    @jwt_required(verify_type=False)                                    # accepts an access or a refresh token
    def logout():
        REVOKED_JTIS.add(get_jwt()["jti"])
        return jsonify(logged_out=True)

    return app


app = create_refresh_app()
client = app.test_client()
tokens = client.post("/api/login", json={"email": "ana@example.com", "password": DEMO_PASSWORD}).get_json()
access = {"Authorization": f"Bearer {tokens['access_token']}"}
refresh = {"Authorization": f"Bearer {tokens['refresh_token']}"}

print("refresh token on a normal endpoint:", client.get("/api/me", headers=refresh).status_code)
print("access token on the refresh endpoint:", client.post("/api/refresh", headers=access).status_code)
rotated = client.post("/api/refresh", headers=refresh).get_json()
print("second use of the same refresh token:", client.post("/api/refresh", headers=refresh).status_code)

stale = {"Authorization": f"Bearer {rotated['access_token']}"}
print("password change with a refreshed token:", client.post("/api/password", headers=stale).status_code)
print("password change with a fresh token:", client.post("/api/password", headers=access).status_code)

client.post("/api/logout", headers=access)
print("after logout:", client.get("/api/me", headers=access).get_json())

from flask_jwt_extended import set_access_cookies, unset_jwt_cookies


def create_cookie_app() -> Flask:
    app = Flask(__name__)
    app.config.update(
        JWT_SECRET_KEY=secrets.token_hex(32),
        JWT_TOKEN_LOCATION=["cookies"],
        JWT_COOKIE_SECURE=False,          # True in production, where the site is served over HTTPS
        JWT_COOKIE_SAMESITE="Strict",
        JWT_COOKIE_CSRF_PROTECT=True,
    )
    JWTManager(app)

    @app.post("/spa/login")
    def login():
        response = jsonify(logged_in=True)
        set_access_cookies(response, create_access_token(identity="1"))   # sets the token and the CSRF cookies
        return response

    @app.post("/spa/logout")
    def logout():
        response = jsonify(logged_out=True)
        unset_jwt_cookies(response)
        return response

    @app.get("/spa/shelves")
    @jwt_required()
    def list_shelves():
        return jsonify(shelves=["to-read"])

    @app.post("/spa/shelves")
    @jwt_required()
    def create_shelf():
        return jsonify(created=True), 201

    return app


app = create_cookie_app()
client = app.test_client()
login_response = client.post("/spa/login")
print([cookie.split("=")[0] for cookie in login_response.headers.getlist("Set-Cookie")])
csrf = client.get_cookie("csrf_access_token").value

print("GET without a header:", client.get("/spa/shelves").status_code)
print("POST without the CSRF header:", client.post("/spa/shelves").get_json())
print("POST with the CSRF header:", client.post("/spa/shelves", headers={"X-CSRF-TOKEN": csrf}).status_code)

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa


def make_rsa_pems() -> tuple[bytes, bytes]:
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    private = key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    )
    public = key.public_key().public_bytes(
        serialization.Encoding.PEM,
        serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    return private, public


private_pem, public_pem = make_rsa_pems()      # in production: generated once and kept in a secret store

issuer_app = Flask("issuer")
issuer_app.config.update(JWT_ALGORITHM="RS256", JWT_PRIVATE_KEY=private_pem, JWT_PUBLIC_KEY=public_pem)
JWTManager(issuer_app)

verifier_app = Flask("verifier")                # a second service: it has no private key at all
verifier_app.config.update(JWT_ALGORITHM="RS256", JWT_PUBLIC_KEY=public_pem)
JWTManager(verifier_app)


@verifier_app.get("/ping")
@jwt_required()
def ping():
    return jsonify(caller=get_jwt_identity())


with issuer_app.app_context():
    issued = create_access_token(identity="1")
print(verifier_app.test_client().get("/ping", headers={"Authorization": f"Bearer {issued}"}).get_json())

with verifier_app.app_context():
    try:
        create_access_token(identity="1")
    except RuntimeError as error:
        print("the verifier cannot issue tokens:", error)

import base64
import hashlib
import hmac
import json
import threading
from functools import wraps
from http.server import BaseHTTPRequestHandler, HTTPServer
import time

import jwt
from flask import abort, g
from jwt.algorithms import RSAAlgorithm

ISSUER = "https://idp.example.test/"
AUDIENCE = "https://api.bookshelf.example.test"

idp_private_pem, idp_public_pem = make_rsa_pems()
idp_private_key = serialization.load_pem_private_key(idp_private_pem, password=None)
jwk = RSAAlgorithm.to_jwk(idp_private_key.public_key(), as_dict=True) | {"kid": "key-1", "use": "sig", "alg": "RS256"}
JWKS_BODY = json.dumps({"keys": [jwk]}).encode()


class JwksHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(JWKS_BODY)

    def log_message(self, *args):           # keep the example output quiet
        pass


jwks_server = HTTPServer(("127.0.0.1", 0), JwksHandler)
threading.Thread(target=jwks_server.serve_forever, daemon=True).start()
JWKS_URL = f"http://127.0.0.1:{jwks_server.server_port}/.well-known/jwks.json"

jwks_client = jwt.PyJWKClient(JWKS_URL, cache_keys=True, lifespan=300)


def verify_access_token(token: str) -> dict:
    signing_key = jwks_client.get_signing_key_from_jwt(token)      # chooses the key by the header's kid
    return jwt.decode(
        token,
        signing_key,
        algorithms=["RS256"],
        audience=AUDIENCE,
        issuer=ISSUER,
        options={"require": ["exp", "iat", "iss", "aud", "sub"]},
        leeway=30,
    )


def external_token_required(scope: str | None = None):
    def decorator(view):
        @wraps(view)
        def wrapped(*args, **kwargs):
            header = request.authorization
            if header is None or header.type != "bearer" or not header.token:
                return jsonify(error="authentication_required"), 401, {"WWW-Authenticate": "Bearer"}
            try:
                g.claims = verify_access_token(header.token)
            except jwt.PyJWTError as error:                        # signature, exp, aud, iss, alg, unknown kid, JWKS down
                return jsonify(error="invalid_token"), 401, {"WWW-Authenticate": 'Bearer error="invalid_token"'}
            if scope and scope not in g.claims.get("scope", "").split():
                abort(403)
            return view(*args, **kwargs)
        return wrapped
    return decorator


resource_app = Flask("resource-server")


@resource_app.get("/api/shelves")
@external_token_required(scope="shelves:read")
def shelves():
    return jsonify(user=g.claims["sub"], audience=g.claims["aud"])


def issue_idp_token(**overrides) -> str:
    now = int(time.time())
    claims = {"iss": ISSUER, "aud": AUDIENCE, "sub": "auth0|ana", "iat": now, "exp": now + 300, "scope": "shelves:read"}
    claims.update(overrides)
    return jwt.encode(claims, idp_private_pem, algorithm="RS256", headers={"kid": "key-1"})


client = resource_app.test_client()


def call(token: str):
    response = client.get("/api/shelves", headers={"Authorization": f"Bearer {token}"})
    return response.status_code, response.get_json()


print("valid:", call(issue_idp_token()))
print("wrong audience:", call(issue_idp_token(aud="https://another-api.example.test"))[0])
print("wrong issuer:", call(issue_idp_token(iss="https://evil.example.test/"))[0])
print("expired:", call(issue_idp_token(exp=int(time.time()) - 3600))[0])
print("missing scope:", call(issue_idp_token(scope="shelves:write"))[0])

def b64(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def forged_hs256_token() -> str:
    """Signs with the PUBLIC key as an HMAC secret (algorithm confusion)."""
    now = int(time.time())
    header = b64(json.dumps({"alg": "HS256", "typ": "JWT", "kid": "key-1"}).encode())
    payload = b64(json.dumps({"iss": ISSUER, "aud": AUDIENCE, "sub": "attacker", "iat": now, "exp": now + 300}).encode())
    signature = hmac.new(idp_public_pem, f"{header}.{payload}".encode(), hashlib.sha256).digest()
    return f"{header}.{payload}.{b64(signature)}"


unsigned = jwt.encode(
    {"iss": ISSUER, "aud": AUDIENCE, "sub": "attacker", "iat": int(time.time()), "exp": int(time.time()) + 300},
    key=None,
    algorithm="none",
    headers={"kid": "key-1"},
)
print("HS256 signed with the public key:", call(forged_hs256_token())[0])
print("alg none:", call(unsigned)[0])
try:
    verify_access_token(forged_hs256_token())
except jwt.PyJWTError as error:
    print(type(error).__name__)

API_KEYS: dict[str, dict] = {}     # a database table in production: prefix, digest, owner, scopes, revoked


def issue_api_key(owner: str, scopes: set[str]) -> str:
    prefix = secrets.token_hex(4)
    secret = secrets.token_urlsafe(32)
    API_KEYS[prefix] = {
        "digest": hashlib.sha256(secret.encode()).digest(),
        "owner": owner,
        "scopes": scopes,
        "revoked": False,
    }
    return f"bk_{prefix}.{secret}"               # returned once and never stored in clear


def api_key_required(scope: str):
    def decorator(view):
        @wraps(view)
        def wrapped(*args, **kwargs):
            presented = request.headers.get("X-API-Key", "")
            prefix, _, secret = presented.removeprefix("bk_").partition(".")
            record = API_KEYS.get(prefix)
            digest = hashlib.sha256(secret.encode()).digest()
            if record is None or record["revoked"] or not hmac.compare_digest(digest, record["digest"]):
                return jsonify(error="invalid_api_key"), 401
            if scope not in record["scopes"]:
                return jsonify(error="insufficient_scope"), 403
            g.api_key_owner = record["owner"]
            return view(*args, **kwargs)
        return wrapped
    return decorator


keys_app = Flask("keys")


@keys_app.get("/partner/books")
@api_key_required("books:read")
def partner_books():
    return jsonify(owner=g.api_key_owner)


reader_key = issue_api_key("acme-sync", {"books:read"})
client = keys_app.test_client()
print(client.get("/partner/books", headers={"X-API-Key": reader_key}).get_json())
print(client.get("/partner/books").status_code, client.get("/partner/books", headers={"X-API-Key": reader_key[:-3] + "xyz"}).status_code)
API_KEYS[reader_key.removeprefix("bk_").split(".")[0]]["revoked"] = True
print("after revocation:", client.get("/partner/books", headers={"X-API-Key": reader_key}).status_code)

from werkzeug.datastructures import Authorization


def basic_required(view):
    @wraps(view)
    def wrapped(*args, **kwargs):
        auth = request.authorization
        user = None
        if auth is not None and auth.type == "basic":
            user = next((u for u in USERS.values() if u["email"] == auth.username), None)
            stored = user["password_hash"] if user else DUMMY_HASH
            if not check_password_hash(stored, auth.password or "") or user is None:
                user = None
        if user is None:
            response = jsonify(error="authentication_required")
            response.status_code = 401
            response.headers["WWW-Authenticate"] = 'Basic realm="Bookshelf", charset="UTF-8"'
            return response
        g.user = user
        return view(*args, **kwargs)
    return wrapped


basic_app = Flask("basic")


@basic_app.get("/internal/stats")
@basic_required
def stats():
    return jsonify(user=g.user["email"], books=3)


client = basic_app.test_client()
anonymous = client.get("/internal/stats")
print(anonymous.status_code, anonymous.headers["WWW-Authenticate"])
print(client.get("/internal/stats", auth=("ana@example.com", DEMO_PASSWORD)).get_json())
print(client.get("/internal/stats", auth=Authorization("basic", {"username": "ana@example.com", "password": secrets.token_urlsafe(8)})).status_code)

from flask_httpauth import HTTPBasicAuth, HTTPTokenAuth

basic = HTTPBasicAuth(realm="Bookshelf")
header_key = HTTPTokenAuth(scheme="Bearer", header="X-API-Key")
shared_key = secrets.token_urlsafe(32)


@basic.verify_password
def verify_password(username, password):
    user = next((u for u in USERS.values() if u["email"] == username), None)
    if user and check_password_hash(user["password_hash"], password):
        return user["email"]
    return None


@header_key.verify_token
def verify_token(token):
    return "robot" if hmac.compare_digest(token, shared_key) else None


httpauth_app = Flask("httpauth")


@httpauth_app.get("/hello")
@basic.login_required
def hello():
    return jsonify(user=basic.current_user())


@httpauth_app.get("/robot")
@header_key.login_required
def robot():
    return jsonify(user=header_key.current_user())


client = httpauth_app.test_client()
print(client.get("/hello").status_code, client.get("/hello", auth=("ana@example.com", DEMO_PASSWORD)).get_json())
print(client.get("/robot", headers={"X-API-Key": shared_key}).get_json(), client.get("/robot").status_code)

from flask import Blueprint
from flask_cors import CORS
from flask_jwt_extended import verify_jwt_in_request

SPA_ORIGIN = "https://app.bookshelf.example.test"


def create_guarded_app(skip_preflight: bool) -> Flask:
    app = Flask(__name__)
    app.config["JWT_SECRET_KEY"] = secrets.token_hex(32)
    JWTManager(app)
    CORS(app, origins=[SPA_ORIGIN], allow_headers=["Authorization", "Content-Type"], max_age=600)

    api = Blueprint("api", __name__, url_prefix="/api")

    @api.before_request
    def require_token():
        if skip_preflight and request.method == "OPTIONS":
            return None                          # the preflight carries no credentials by design
        if "Authorization" not in request.headers:
            return jsonify(error="authentication_required"), 401
        verify_jwt_in_request()                  # raises the error that JWTManager turns into a 401

    @api.get("/books")
    def books():
        return jsonify(books=[], user=get_jwt_identity())

    app.register_blueprint(api)
    return app


preflight_headers = {
    "Origin": SPA_ORIGIN,
    "Access-Control-Request-Method": "GET",
    "Access-Control-Request-Headers": "authorization",
}
broken = create_guarded_app(skip_preflight=False).test_client().options("/api/books", headers=preflight_headers)
fixed = create_guarded_app(skip_preflight=True).test_client().options("/api/books", headers=preflight_headers)
print("guard rejects the preflight:", broken.status_code)
print("guard skips the preflight:", fixed.status_code, fixed.headers["Access-Control-Allow-Origin"], fixed.headers["Access-Control-Allow-Headers"])

SHELVES = {17: {"id": 17, "owner": "1", "name": "to-read"}, 18: {"id": 18, "owner": "2", "name": "private"}}

owner_app = create_json_error_app()


@owner_app.get("/api/shelves/<int:shelf_id>")
@jwt_required()
def get_shelf(shelf_id: int):
    shelf = SHELVES.get(shelf_id)
    if shelf is None or shelf["owner"] != get_jwt_identity():
        abort(404)                                # 404, not 403: do not reveal that the shelf exists
    return jsonify(shelf)


client = owner_app.test_client()
tokens = client.post("/api/login", json={"email": "ana@example.com", "password": DEMO_PASSWORD}).get_json()
bearer = {"Authorization": f"Bearer {tokens['access_token']}"}
print(client.get("/api/shelves/17", headers=bearer).status_code, client.get("/api/shelves/18", headers=bearer).status_code)

from flask_smorest import Api, Blueprint as SmorestBlueprint

doc_app = create_json_error_app()
doc_app.config.update(
    API_TITLE="Bookshelf API",
    API_VERSION="v1",
    OPENAPI_VERSION="3.1.0",
    OPENAPI_URL_PREFIX="/",
    OPENAPI_JSON_PATH="openapi.json",
    OPENAPI_SWAGGER_UI_PATH="/docs",
    OPENAPI_SWAGGER_UI_URL="https://cdn.jsdelivr.net/npm/swagger-ui-dist/",
)
api = Api(doc_app)
api.spec.components.security_scheme("bearerAuth", {"type": "http", "scheme": "bearer", "bearerFormat": "JWT"})

shelves_blp = SmorestBlueprint("shelves", __name__, url_prefix="/v1/shelves", description="Shelves")


@shelves_blp.route("/")
@shelves_blp.doc(security=[{"bearerAuth": []}])
@jwt_required()
def list_my_shelves():
    return jsonify([s for s in SHELVES.values() if s["owner"] == get_jwt_identity()])


api.register_blueprint(shelves_blp)

spec = doc_app.test_client().get("/openapi.json").get_json()
print(spec["components"]["securitySchemes"])
print(spec["paths"]["/v1/shelves/"]["get"]["security"])