Update management command for creating jwks

openedx/core/djangoapps/oauth_dispatch/management/commands/generate_jwt_signing_key.py

@@ -14,7 +14,7 @@
 from Cryptodome.PublicKey import RSA
 from django.conf import settings
 from django.core.management.base import BaseCommand
-from jwkest import jwk
+from jwt.algorithms import get_default_algorithms
 
 log = logging.getLogger(__name__)
 
@@ -123,15 +123,23 @@ def _generate_key_id(self, size, chars=string.ascii_uppercase + string.digits):
     def _generate_key_pair(self, key_size, key_id):
         log.info('Generating new JWT signing keypair for key id %s.', key_id)
         rsa_key = RSA.generate(key_size)
-        rsa_jwk = jwk.RSAKey(kid=key_id, key=rsa_key)
-        return rsa_jwk
+        algo = get_default_algorithms()['RS512']
+        key_data = algo.prepare_key(rsa_key.export_key('PEM').decode())
+        rsa_jwk = json.loads(algo.to_jwk(key_data))
+        public_rsa_jwk = json.loads(algo.to_jwk(key_data.public_key()))
+
+        rsa_jwk['kid'] = key_id
+        public_rsa_jwk['kid'] = key_id
+        return {'private': rsa_jwk, 'public': public_rsa_jwk}
 
     def _output_public_keys(self, jwk_key, add_previous, strip_prefix):
-        public_keys = jwk.KEYS()
+        public_keys = {'keys': []}
+
         if add_previous:
             self._add_previous_public_keys(public_keys)
-        public_keys.append(jwk_key)
-        serialized_public_keys = public_keys.dump_jwks()
+
+        public_keys['keys'].append(jwk_key['public'])
+        serialized_public_keys = json.dumps(public_keys)
 
         prefix = '' if strip_prefix else 'COMMON_'
         public_signing_key = f'{prefix}JWT_PUBLIC_SIGNING_JWK_SET'
@@ -155,11 +163,10 @@ def _add_previous_public_keys(self, public_keys):
         previous_signing_keys = settings.JWT_AUTH.get('JWT_PUBLIC_SIGNING_JWK_SET')
         if previous_signing_keys:
             log.info('Old JWT_PUBLIC_SIGNING_JWK_SET: %s.', previous_signing_keys)
-            public_keys.load_jwks(previous_signing_keys)
+            public_keys['keys'].extend(json.loads(previous_signing_keys)['keys'])
 
     def _output_private_keys(self, jwk_key, strip_prefix):
-        serialized_keypair = jwk_key.serialize(private=True)
-        serialized_keypair_json = json.dumps(serialized_keypair)
+        serialized_keypair_json = json.dumps(jwk_key['private'])
 
         prefix = '' if strip_prefix else 'EDXAPP_'
         private_signing_key = f'{prefix}JWT_PRIVATE_SIGNING_JWK'

requirements/edx/base.in

@@ -147,8 +147,6 @@ py2neo                              # Driver for converting Python modulestore s
 pycountry
 pycryptodomex
 pygments                            # Used to support colors in paver command output
-pyjwkest
-# TODO Replace PyJWT usage with pyjwkest
 # PyJWT 1.6.3 contains PyJWTError, which is required by Apple auth in social-auth-core
 PyJWT>=1.6.3
 pylti1p3                            # Required by content_libraries core library to suport LTI 1.3 launches

requirements/edx/base.txt

@@ -859,7 +859,6 @@ pygments==2.15.1
     #   py2neo
 pyjwkest==1.4.2
     # via
-    #   -r requirements/edx/base.in
     #   edx-token-utils
     #   lti-consumer-xblock
 pyjwt[crypto]==2.7.0