PyJWT — Security Pitfalls & Hardening¶

# test_pitfalls.py
import base64
import hashlib
import hmac
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import jwt
import pytest
from cryptography.hazmat.primitives import serialization as ser
from cryptography.hazmat.primitives.asymmetric import rsa

SECRET = "x" * 32                        # placeholder: 32 bytes for HS256


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


@pytest.fixture(scope="module")
def rsa_pair():
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    private = key.private_bytes(ser.Encoding.PEM, ser.PrivateFormat.PKCS8, ser.NoEncryption())
    public = key.public_key().public_bytes(ser.Encoding.PEM, ser.PublicFormat.SubjectPublicKeyInfo)
    return private, public


def test_payload_is_readable_without_the_key():
    token = jwt.encode({"sub": "u1", "email": "[email protected]"}, SECRET, algorithm="HS256")
    claims = jwt.decode(token, options={"verify_signature": False})
    assert claims["email"] == "[email protected]"


def test_algorithms_is_required():
    token = jwt.encode({"sub": "u1"}, SECRET, algorithm="HS256")
    with pytest.raises(jwt.DecodeError, match="algorithms"):
        jwt.decode(token, SECRET)


def test_algorithms_must_be_a_list():
    token = jwt.encode({"sub": "u1"}, SECRET, algorithm="HS256")
    # a bare string is matched as a substring: this is NOT an allow-list
    assert jwt.decode(token, SECRET, algorithms="HS256 RS256") == {"sub": "u1"}

# test_pitfalls.py (continued)
def test_alg_none_is_rejected():
    unsigned = jwt.encode({"sub": "admin"}, None, algorithm="none")
    with pytest.raises(jwt.InvalidAlgorithmError):
        jwt.decode(unsigned, SECRET, algorithms=["HS256"])
    with pytest.raises(jwt.InvalidSignatureError):       # even when "none" is allowed
        jwt.decode(unsigned, None, algorithms=["none"])

# test_pitfalls.py (continued)
def forge_hs256_with_public_key(public_pem: bytes) -> str:
    header = b64(json.dumps({"alg": "HS256", "typ": "JWT"}).encode())
    body = b64(json.dumps({"sub": "admin"}).encode())
    mac = hmac.new(public_pem, f"{header}.{body}".encode(), hashlib.sha256)
    return f"{header}.{body}.{b64(mac.digest())}"


def test_algorithm_confusion_is_stopped(rsa_pair):
    _, public = rsa_pair
    forged = forge_hs256_with_public_key(public)
    with pytest.raises(jwt.InvalidAlgorithmError):                       # the allow-list
        jwt.decode(forged, public, algorithms=["RS256"])
    with pytest.raises(jwt.InvalidKeyError, match="asymmetric key"):     # PyJWT's key-type check
        jwt.decode(forged, public, algorithms=["RS256", "HS256"])
    with pytest.raises(jwt.InvalidKeyError):                             # it will not sign either
        jwt.encode({"sub": "admin"}, public, algorithm="HS256")

# test_pitfalls.py (continued)
class RevocationList:
    """In-memory sketch. In production use Redis (SET jti 1 EX ttl) or a DB table."""

    def __init__(self) -> None:
        self._revoked: dict[str, float] = {}

    def revoke(self, claims: dict) -> None:
        self._revoked[claims["jti"]] = claims["exp"]

    def is_revoked(self, jti: str) -> bool:
        now = time.time()
        self._revoked = {j: exp for j, exp in self._revoked.items() if exp > now}   # purge
        return jti in self._revoked


def test_revocation_by_jti():
    token = jwt.encode({"jti": "abc", "exp": int(time.time()) + 60}, SECRET, algorithm="HS256")
    claims = jwt.decode(token, SECRET, algorithms=["HS256"], options={"require": ["jti", "exp"]})
    revoked = RevocationList()
    assert not revoked.is_revoked(claims["jti"])
    revoked.revoke(claims)
    assert revoked.is_revoked(claims["jti"])

from fastapi import FastAPI, Response

app = FastAPI()


@app.post("/auth/login")
def login(response: Response):
    response.set_cookie(
        "refresh_token", "<token>",
        httponly=True, secure=True, samesite="strict", max_age=7 * 24 * 3600, path="/auth",
    )
    return {"access_token": "<token>", "token_type": "bearer"}

# test_pitfalls.py (continued)
def test_short_hmac_key_warns():
    # pytest.ini promotes the warning to an error:
    #   [pytest]
    #   filterwarnings = error::jwt.warnings.InsecureKeyLengthWarning
    with pytest.raises(jwt.InsecureKeyLengthWarning):
        jwt.encode({"sub": "u1"}, "short", algorithm="HS256")
    strict = jwt.PyJWT({"enforce_minimum_key_length": True})
    with pytest.raises(jwt.InvalidKeyError):
        strict.encode({"sub": "u1"}, "short", algorithm="HS256")

# test_pitfalls.py (continued)
KEYS = {"2026-10": SECRET}


def resolve_key(token: str) -> str:
    kid = jwt.get_unverified_header(token).get("kid")     # kid must be a str, else InvalidTokenError
    try:
        return KEYS[kid]
    except KeyError:
        raise jwt.InvalidTokenError("unknown kid") from None


def test_kid_is_only_a_lookup_key():
    good = jwt.encode({"sub": "u1"}, SECRET, algorithm="HS256", headers={"kid": "2026-10"})
    assert jwt.decode(good, resolve_key(good), algorithms=["HS256"]) == {"sub": "u1"}
    for kid in ("../../etc/passwd", "1' OR '1'='1"):
        bad = jwt.encode({"sub": "u1"}, SECRET, algorithm="HS256", headers={"kid": kid})
        with pytest.raises(jwt.InvalidTokenError):
            resolve_key(bad)
    header = b64(json.dumps({"alg": "HS256", "kid": 5}).encode())     # hand-built: encode() refuses it
    with pytest.raises(jwt.InvalidTokenError):
        resolve_key(f"{header}.{b64(b'{}')}.{b64(b'sig')}")


def test_unknown_crit_extension_is_rejected():
    token = jwt.encode({"sub": "u1"}, SECRET, algorithm="HS256",
                       headers={"crit": ["exp-ext"], "exp-ext": 1})
    with pytest.raises(jwt.InvalidTokenError, match="Unsupported critical extension"):
        jwt.decode(token, SECRET, algorithms=["HS256"])


def test_jku_header_is_never_fetched():
    token = jwt.encode({"sub": "u1"}, SECRET, algorithm="HS256",
                       headers={"jku": "https://attacker.example/jwks.json"})
    assert jwt.decode(token, SECRET, algorithms=["HS256"]) == {"sub": "u1"}


def test_jwks_client_rejects_other_schemes_and_redirects():
    with pytest.raises(jwt.PyJWKClientError, match="scheme"):
        jwt.PyJWKClient("file:///etc/passwd")

    class Redirect(BaseHTTPRequestHandler):
        def do_GET(self):
            self.send_response(302)
            self.send_header("Location", "http://127.0.0.1:1/jwks.json")
            self.end_headers()

        def log_message(self, *args):
            pass

    server = ThreadingHTTPServer(("127.0.0.1", 0), Redirect)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    try:
        client = jwt.PyJWKClient(f"http://127.0.0.1:{server.server_port}/jwks.json", timeout=2)
        with pytest.raises(jwt.PyJWKClientConnectionError):
            client.get_signing_key("any")
    finally:
        server.shutdown()
        server.server_close()