Use PyJWT+cryptography instead of jwkest+Cryptodrome (juanifioren#446)

- Switches id_token signing and verification to use PyJWT with cryptography
- Removes jwkest and Cryptodome dependencies
- `future` dependency previously required by jwkest is no longer needed (it had unfixed security vulnerabilities and seems unmaintained)
- adds caching of RSA keys to avoid repeated expensive key loading operations

Author: stefanfoulis

oidc_provider/lib/utils/token.py

@@ -2,13 +2,10 @@
 import uuid
 from datetime import timedelta
 
-from Cryptodome.PublicKey.RSA import importKey
+import jwt
+from cryptography.hazmat.primitives import serialization
 from django.utils import dateformat
 from django.utils import timezone
-from jwkest.jwk import RSAKey as jwk_RSAKey
-from jwkest.jwk import SYMKey
-from jwkest.jws import JWS
-from jwkest.jwt import JWT
 
 from oidc_provider import settings
 from oidc_provider.lib.claims import StandardScopeClaims
@@ -18,6 +15,10 @@
 from oidc_provider.models import RSAKey
 from oidc_provider.models import Token
 
+# Cache for loaded RSA keys to avoid repeated PEM parsing
+# Cache is automatically cleaned of stale entries (keys no longer in DB)
+_rsa_key_cache = {}
+
 
 def create_id_token(token, user, aud, nonce="", at_hash="", request=None, scope=None):
     """
@@ -72,28 +73,56 @@ def create_id_token(token, user, aud, nonce="", at_hash="", request=None, scope=
 def encode_id_token(payload, client):
     """
     Represent the ID Token as a JSON Web Token (JWT).
-    Return a hash.
+    Returns a dict.
     """
     keys = get_client_alg_keys(client)
-    _jws = JWS(payload, alg=client.jwt_alg)
-    return _jws.sign_compact(keys)
+    # Use the first key for encoding
+    # TODO: make key selection more explicit
+    key_info = keys[0]
+
+    headers = {}
+    if "kid" in key_info:
+        headers["kid"] = key_info["kid"]
+
+    return jwt.encode(payload, key_info["key"], algorithm=key_info["algorithm"], headers=headers)
 
 
 def decode_id_token(token, client):
     """
     Represent the ID Token as a JSON Web Token (JWT).
-    Return a hash.
+    Returns a dict.
     """
-    keys = get_client_alg_keys(client)
-    return JWS().verify_compact(token, keys=keys)
+    # Try decoding with each available key
+    for key in get_client_alg_keys(client):
+        try:
+            return jwt.decode(
+                jwt=token,
+                # HS256 uses the same key for signing and verifying
+                key=key["key"] if key["algorithm"] == "HS256" else key["public_key"],
+                algorithms=[key["algorithm"]],
+                options={
+                    "verify_signature": True,
+                    "verify_aud": False,  # Disable audience validation for compatibility
+                    "verify_exp": False,  # Disable expiration validation for compatibility
+                    "verify_iat": False,  # Disable issued at validation for compatibility
+                    "verify_nbf": False,  # Disable not before validation for compatibility
+                },
+            )
+        except jwt.InvalidTokenError:
+            continue
+
+    # If we get here, none of the keys worked
+    raise jwt.InvalidTokenError("Token could not be decoded with any available key")
 
 
 def client_id_from_id_token(id_token):
     """
     Extracts the client id from a JSON Web Token (JWT).
+    Does NOT verify the token signature or expiration.
     Returns a string or None.
     """
-    payload = JWT().unpack(id_token).payload()
+    # Decode without verification to get the payload
+    payload = jwt.decode(id_token, options={"verify_signature": False})
     aud = payload.get("aud", None)
     if aud is None:
         return None
@@ -150,16 +179,47 @@ def create_code(
 def get_client_alg_keys(client):
     """
     Takes a client and returns the set of keys associated with it.
-    Returns a list of keys.
+    Returns a list of keys compatible with PyJWT.
     """
     if client.jwt_alg == "RS256":
         keys = []
+        current_kids = set()
+
         for rsakey in RSAKey.objects.all():
-            keys.append(jwk_RSAKey(key=importKey(rsakey.key), kid=rsakey.kid))
+            cache_key = f"rsa_key_{rsakey.kid}"
+            current_kids.add(cache_key)
+
+            if cache_key not in _rsa_key_cache:
+                # Load the RSA private key using cryptography (expensive operation)
+                private_key = serialization.load_pem_private_key(
+                    rsakey.key.encode("utf-8"),
+                    password=None,
+                )
+                # Also cache the public key to avoid repeated .public_key() calls
+                public_key = private_key.public_key()
+                _rsa_key_cache[cache_key] = {"private_key": private_key, "public_key": public_key}
+
+            key_pair = _rsa_key_cache[cache_key]
+            keys.append(
+                {
+                    "key": key_pair["private_key"],
+                    "public_key": key_pair["public_key"],
+                    "kid": rsakey.kid,
+                    "algorithm": "RS256",
+                }
+            )
+
+        # Clean up stale cache entries (keys that no longer exist in DB)
+        stale_keys = set(_rsa_key_cache.keys()) - current_kids
+        for stale_key in stale_keys:
+            del _rsa_key_cache[stale_key]
+
         if not keys:
             raise Exception("You must add at least one RSA Key.")
     elif client.jwt_alg == "HS256":
-        keys = [SYMKey(key=client.client_secret, alg=client.jwt_alg)]
+        # NOTE: HS256 does not have any expensive key parsing, so we don't need the
+        #       same key caching as RS256.
+        keys = [{"key": client.client_secret, "algorithm": "HS256"}]
     else:
         raise Exception("Unsupported key algorithm.")
 

oidc_provider/management/commands/creatersakey.py

@@ -1,4 +1,5 @@
-from Cryptodome.PublicKey import RSA
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 from django.core.management.base import BaseCommand
 
 from oidc_provider.models import RSAKey
@@ -9,8 +10,20 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         try:
-            key = RSA.generate(2048)
-            rsakey = RSAKey(key=key.exportKey("PEM").decode("utf8"))
+            # Generate a new RSA private key with 2048 bits
+            private_key = rsa.generate_private_key(
+                public_exponent=65537,
+                key_size=2048,
+            )
+
+            # Serialize the private key to PEM format
+            key_pem = private_key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            ).decode("utf-8")
+
+            rsakey = RSAKey(key=key_pem)
             rsakey.save()
             self.stdout.write("RSA key successfully created with kid: {0}".format(rsakey.kid))
         except Exception as e:

oidc_provider/tests/cases/test_authorize_endpoint.py

@@ -22,12 +22,12 @@
     from django.urls import reverse
 except ImportError:
     from django.core.urlresolvers import reverse
+import jwt
 from django.contrib.auth.models import AnonymousUser
 from django.core.management import call_command
 from django.test import RequestFactory
 from django.test import TestCase
 from django.test import override_settings
-from jwkest.jwt import JWT
 
 from oidc_provider import settings
 from oidc_provider.lib.endpoints.authorize import AuthorizeEndpoint
@@ -724,7 +724,7 @@ def test_idtoken_token_at_hash(self):
         # obtain `id_token` portion of Location
         components = urlsplit(response["Location"])
         fragment = parse_qs(components[4])
-        id_token = JWT().unpack(fragment["id_token"][0].encode("utf-8")).payload()
+        id_token = jwt.decode(fragment["id_token"][0], options={"verify_signature": False})
 
         self.assertIn("at_hash", id_token)
 
@@ -750,7 +750,7 @@ def test_idtoken_at_hash(self):
         # obtain `id_token` portion of Location
         components = urlsplit(response["Location"])
         fragment = parse_qs(components[4])
-        id_token = JWT().unpack(fragment["id_token"][0].encode("utf-8")).payload()
+        id_token = jwt.decode(fragment["id_token"][0], options={"verify_signature": False})
 
         self.assertNotIn("at_hash", id_token)