feat!: remove pyjwkest (#36707)

Co-authored-by: M Umar Khan <[email protected]>

Author: mumarkhan999

lms/envs/test.py

@@ -651,28 +651,40 @@
     'JWT_ISSUER': 'token-test-issuer',
     'JWT_SIGNING_ALGORITHM': 'RS512',
     'JWT_SUPPORTED_VERSION': '1.2.0',
-    'JWT_PRIVATE_SIGNING_JWK': '''{
-        "e": "AQAB",
-        "d": "HIiV7KNjcdhVbpn3KT-I9n3JPf5YbGXsCIedmPqDH1d4QhBofuAqZ9zebQuxkRUpmqtYMv0Zi6ECSUqH387GYQF_XvFUFcjQRPycISd8TH0DAKaDpGr-AYNshnKiEtQpINhcP44I1AYNPCwyoxXA1fGTtmkKChsuWea7o8kytwU5xSejvh5-jiqu2SF4GEl0BEXIAPZsgbzoPIWNxgO4_RzNnWs6nJZeszcaDD0CyezVSuH9QcI6g5QFzAC_YuykSsaaFJhZ05DocBsLczShJ9Omf6PnK9xlm26I84xrEh_7x4fVmNBg3xWTLh8qOnHqGko93A1diLRCrKHOvnpvgQ",
-        "n": "o5cn3ljSRi6FaDEKTn0PS-oL9EFyv1pI7dRgffQLD1qf5D6sprmYfWWokSsrWig8u2y0HChSygR6Jn5KXBqQn6FpM0dDJLnWQDRXHLl3Ey1iPYgDSmOIsIGrV9ZyNCQwk03wAgWbfdBTig3QSDYD-sTNOs3pc4UD_PqAvU2nz_1SS2ZiOwOn5F6gulE1L0iE3KEUEvOIagfHNVhz0oxa_VRZILkzV-zr6R_TW1m97h4H8jXl_VJyQGyhMGGypuDrQ9_vaY_RLEulLCyY0INglHWQ7pckxBtI5q55-Vio2wgewe2_qYcGsnBGaDNbySAsvYcWRrqDiFyzrJYivodqTQ",
-        "q": "3T3DEtBUka7hLGdIsDlC96Uadx_q_E4Vb1cxx_4Ss_wGp1Loz3N3ZngGyInsKlmbBgLo1Ykd6T9TRvRNEWEtFSOcm2INIBoVoXk7W5RuPa8Cgq2tjQj9ziGQ08JMejrPlj3Q1wmALJr5VTfvSYBu0WkljhKNCy1KB6fCby0C9WE",
-        "p": "vUqzWPZnDG4IXyo-k5F0bHV0BNL_pVhQoLW7eyFHnw74IOEfSbdsMspNcPSFIrtgPsn7981qv3lN_staZ6JflKfHayjB_lvltHyZxfl0dvruShZOx1N6ykEo7YrAskC_qxUyrIvqmJ64zPW3jkuOYrFs7Ykj3zFx3Zq1H5568G0",
-        "kid": "token-test-sign", "kty": "RSA"
-    }''',
-    'JWT_PUBLIC_SIGNING_JWK_SET': '''{
-        "keys": [
+    'JWT_PRIVATE_SIGNING_JWK': """
+        {
+            "kid": "token-test-sign",
+            "kty": "RSA",
+            "key_ops": [
+                "sign"
+            ],
+            "n": "o5cn3ljSRi6FaDEKTn0PS-oL9EFyv1pI7dRgffQLD1qf5D6sprmYfWWokSsrWig8u2y0HChSygR6Jn5KXBqQn6FpM0dDJLnWQDRXHLl3Ey1iPYgDSmOIsIGrV9ZyNCQwk03wAgWbfdBTig3QSDYD-sTNOs3pc4UD_PqAvU2nz_1SS2ZiOwOn5F6gulE1L0iE3KEUEvOIagfHNVhz0oxa_VRZILkzV-zr6R_TW1m97h4H8jXl_VJyQGyhMGGypuDrQ9_vaY_RLEulLCyY0INglHWQ7pckxBtI5q55-Vio2wgewe2_qYcGsnBGaDNbySAsvYcWRrqDiFyzrJYivodqTQ",
+            "e": "AQAB",
+            "d": "HIiV7KNjcdhVbpn3KT-I9n3JPf5YbGXsCIedmPqDH1d4QhBofuAqZ9zebQuxkRUpmqtYMv0Zi6ECSUqH387GYQF_XvFUFcjQRPycISd8TH0DAKaDpGr-AYNshnKiEtQpINhcP44I1AYNPCwyoxXA1fGTtmkKChsuWea7o8kytwU5xSejvh5-jiqu2SF4GEl0BEXIAPZsgbzoPIWNxgO4_RzNnWs6nJZeszcaDD0CyezVSuH9QcI6g5QFzAC_YuykSsaaFJhZ05DocBsLczShJ9Omf6PnK9xlm26I84xrEh_7x4fVmNBg3xWTLh8qOnHqGko93A1diLRCrKHOvnpvgQ",
+            "p": "3T3DEtBUka7hLGdIsDlC96Uadx_q_E4Vb1cxx_4Ss_wGp1Loz3N3ZngGyInsKlmbBgLo1Ykd6T9TRvRNEWEtFSOcm2INIBoVoXk7W5RuPa8Cgq2tjQj9ziGQ08JMejrPlj3Q1wmALJr5VTfvSYBu0WkljhKNCy1KB6fCby0C9WE",
+            "q": "vUqzWPZnDG4IXyo-k5F0bHV0BNL_pVhQoLW7eyFHnw74IOEfSbdsMspNcPSFIrtgPsn7981qv3lN_staZ6JflKfHayjB_lvltHyZxfl0dvruShZOx1N6ykEo7YrAskC_qxUyrIvqmJ64zPW3jkuOYrFs7Ykj3zFx3Zq1H5568G0",
+            "dp": "Azh08H8r2_sJuBXAzx_mQ6iZnAZQ619PnJFOXjTqnMgcaK8iSHLL2CgDIUQwteUcBphgP0uBrfWIBs5jmM8rUtVz4CcrPb5jdjhHjuu4NxmnFbPlhNoOp8OBUjPP3S-h-fPoaFjxDrUqz_zCdPVzp4S6UTkf6Hu-SiI9CFVFZ8E",
+            "dq": "WQ44_KTIbIej9qnYUPMA1DoaAF8ImVDIdiOp9c79dC7FvCpN3w-lnuugrYDM1j9Tk5bRrY7-JuE6OaKQgOtajoS1BIxjYHj5xAVPD15CVevOihqeq5Zx0ZAAYmmCKRrfUe0iLx2QnIcoKH1-Azs23OXeeo6nysznZjvv9NVJv60",
+            "qi": "KSWGH607H1kNG2okjYdmVdNgLxTUB-Wye9a9FNFE49UmQIOJeZYXtDzcjk8IiK3g-EU3CqBeDKVUgHvHFu4_Wj3IrIhKYizS4BeFmOcPDvylDQCmJcC9tXLQgHkxM_MEJ7iLn9FOLRshh7GPgZphXxMhezM26Cz-8r3_mACHu84"
+        }
+    """,  # noqa: E501,
+
+    'JWT_PUBLIC_SIGNING_JWK_SET': """
+        {
+          "keys": [
             {
-                "kid":"token-test-wrong-key",
-                "e": "AQAB",
-                "kty": "RSA",
-                "n": "o5cn3ljSRi6FaDEKTn0PS-oL9EFyv1pI7dffgRQLD1qf5D6sprmYfWVokSsrWig8u2y0HChSygR6Jn5KXBqQn6FpM0dDJLnWQDRXHLl3Ey1iPYgDSmOIsIGrV9ZyNCQwk03wAgWbfdBTig3QSDYD-sTNOs3pc4UD_PqAvU2nz_1SS2ZiOwOn5F6gulE1L0iE3KEUEvOIagfHNVhz0oxa_VRZILkzV-zr6R_TW1m97h4H8jXl_VJyQGyhMGGypuDrQ9_vaY_RLEulLCyY0INglHWQ7pckxBtI5q55-Vio2wgewe2_qYcGsnBGaDNbySAsvYcWRrqDiFyzrJYivodqTQ"
+              "kid": "token-test-sign",
+              "kty": "RSA",
+              "n": "o5cn3ljSRi6FaDEKTn0PS-oL9EFyv1pI7dRgffQLD1qf5D6sprmYfWWokSsrWig8u2y0HChSygR6Jn5KXBqQn6FpM0dDJLnWQDRXHLl3Ey1iPYgDSmOIsIGrV9ZyNCQwk03wAgWbfdBTig3QSDYD-sTNOs3pc4UD_PqAvU2nz_1SS2ZiOwOn5F6gulE1L0iE3KEUEvOIagfHNVhz0oxa_VRZILkzV-zr6R_TW1m97h4H8jXl_VJyQGyhMGGypuDrQ9_vaY_RLEulLCyY0INglHWQ7pckxBtI5q55-Vio2wgewe2_qYcGsnBGaDNbySAsvYcWRrqDiFyzrJYivodqTQ",
+              "e": "AQAB"
             },
             {
-                "kid":"token-test-sign",
-                "e": "AQAB",
-                "kty": "RSA",
-                "n": "o5cn3ljSRi6FaDEKTn0PS-oL9EFyv1pI7dRgffQLD1qf5D6sprmYfWWokSsrWig8u2y0HChSygR6Jn5KXBqQn6FpM0dDJLnWQDRXHLl3Ey1iPYgDSmOIsIGrV9ZyNCQwk03wAgWbfdBTig3QSDYD-sTNOs3pc4UD_PqAvU2nz_1SS2ZiOwOn5F6gulE1L0iE3KEUEvOIagfHNVhz0oxa_VRZILkzV-zr6R_TW1m97h4H8jXl_VJyQGyhMGGypuDrQ9_vaY_RLEulLCyY0INglHWQ7pckxBtI5q55-Vio2wgewe2_qYcGsnBGaDNbySAsvYcWRrqDiFyzrJYivodqTQ"
+              "kid": "token-test-wrong-key",
+              "kty": "RSA",
+              "n": "o5cn3ljSRi6FaDEKTn0PS-oL9EFyv1pI7dRgffQLD1qf5D6sprmYfWWokSsrWig8u2y0HChSygR6Jn5KXBqQn6FpM0dDJLnWQDRXHLl3Ey1iPYgDSmOIsIGrV9ZyNCQwk03wAgWbfdBTig3QSDYD-sTNOs3pc4UD_PqAvU2nz_1SS2ZiOwOn5F6gulE1L0iE3KEUEvOIagfHNVhz0oxa_VRZILkzV-zr6R_TW1m97h4H8jXl_VJyQGyhMGGypuDrQ9_vaY_RLEulLCyY0INglHWQ7pckxBtI5q55-Vio2wgewe2_qYcGsnBGaDNbySAsvYcWRrqDiFyzrJYivodqTQ",
+              "e": "AQAB"
             }
-        ]
-    }''',
+          ]
+        }
+    """,  # noqa: E501
 }

openedx/core/lib/jwt.py

@@ -2,12 +2,12 @@
 JWT Token handling and signing functions.
 """
 
-import json
+import jwt
 from time import time
 
 from django.conf import settings
-from jwkest import Expired, Invalid, MissingKey, jwk
-from jwkest.jws import JWS
+from jwt.api_jwk import PyJWK, PyJWKSet
+from jwt.exceptions import ExpiredSignatureError, InvalidSignatureError, MissingRequiredClaimError
 
 
 def create_jwt(lms_user_id, expires_in_seconds, additional_token_claims, now=None):
@@ -40,15 +40,9 @@ def _encode_and_sign(payload):
 
     The signing key and algorithm are pulled from settings.
     """
-    keys = jwk.KEYS()
-
-    serialized_keypair = json.loads(settings.TOKEN_SIGNING['JWT_PRIVATE_SIGNING_JWK'])
-    keys.add(serialized_keypair)
+    private_key = PyJWK.from_json(settings.TOKEN_SIGNING['JWT_PRIVATE_SIGNING_JWK'])
     algorithm = settings.TOKEN_SIGNING['JWT_SIGNING_ALGORITHM']
-
-    data = json.dumps(payload)
-    jws = JWS(data, alg=algorithm)
-    return jws.sign_compact(keys=keys)
+    return jwt.encode(payload, key=private_key.key, algorithm=algorithm)
 
 
 def unpack_jwt(token, lms_user_id, now=None):
@@ -65,27 +59,40 @@ def unpack_jwt(token, lms_user_id, now=None):
     Returns a valid, decoded json payload (string).
     """
     now = now or int(time())
-    payload = _unpack_and_verify(token)
+    payload = unpack_and_verify(token)
 
     if "lms_user_id" not in payload:
-        raise MissingKey("LMS user id is missing")
+        raise MissingRequiredClaimError("LMS user id is missing")
     if "exp" not in payload:
-        raise MissingKey("Expiration is missing")
+        raise MissingRequiredClaimError("Expiration is missing")
     if payload["lms_user_id"] != lms_user_id:
-        raise Invalid("User does not match")
+        raise InvalidSignatureError("User does not match")
     if payload["exp"] < now:
-        raise Expired("Token is expired")
+        raise ExpiredSignatureError("Token is expired")
 
     return payload
 
 
-def _unpack_and_verify(token):
+def unpack_and_verify(token):  # pylint: disable=inconsistent-return-statements
     """
     Unpack and verify the provided token.
 
     The signing key and algorithm are pulled from settings.
     """
-    keys = jwk.KEYS()
-    keys.load_jwks(settings.TOKEN_SIGNING['JWT_PUBLIC_SIGNING_JWK_SET'])
-    decoded = JWS().verify_compact(token.encode('utf-8'), keys)
-    return decoded
+    key_set = []
+    key_set.extend(
+        PyJWKSet.from_json(settings.TOKEN_SIGNING["JWT_PUBLIC_SIGNING_JWK_SET"]).keys
+    )
+
+    for i in range(len(key_set)):  # pylint: disable=consider-using-enumerate
+        try:
+            decoded = jwt.decode(
+                token,
+                key=key_set[i].key,
+                algorithms=["RS256", "RS512"],
+                options={"verify_signature": True, "verify_aud": False},
+            )
+            return decoded
+        except Exception:  # pylint: disable=broad-exception-caught
+            if i == len(key_set) - 1:
+                raise

openedx/core/lib/tests/test_jwt.py

@@ -2,24 +2,23 @@
 Tests for token handling
 """
 import unittest
+from time import time
 
-from django.conf import settings
-from jwkest import BadSignature, Expired, Invalid, MissingKey, jwk
-from jwkest.jws import JWS
+from jwt.exceptions import ExpiredSignatureError, InvalidSignatureError, MissingRequiredClaimError
 
 from openedx.core.djangolib.testing.utils import skip_unless_lms
-from openedx.core.lib.jwt import _encode_and_sign, create_jwt, unpack_jwt
+from openedx.core.lib.jwt import _encode_and_sign, create_jwt, unpack_jwt, unpack_and_verify
 
 
 test_user_id = 121
 invalid_test_user_id = 120
-test_timeout = 60
-test_now = 1661432902
+test_timeout = 1000
+test_now = int(time())
 test_claims = {"foo": "bar", "baz": "quux", "meaning": 42}
 expected_full_token = {
     "lms_user_id": test_user_id,
-    "iat": 1661432902,
-    "exp": 1661432902 + 60,
+    "iat": test_now,
+    "exp": test_now + test_timeout,
     "iss": "token-test-issuer",  # these lines from test_settings.py
     "version": "1.2.0",  # these lines from test_settings.py
 }
@@ -34,7 +33,7 @@ class TestSign(unittest.TestCase):
     def test_create_jwt(self):
         token = create_jwt(test_user_id, test_timeout, {}, test_now)
 
-        decoded = _verify_jwt(token)
+        decoded = unpack_and_verify(token)
         self.assertEqual(expected_full_token, decoded)
 
     def test_create_jwt_with_claims(self):
@@ -43,7 +42,7 @@ def test_create_jwt_with_claims(self):
         expected_token_with_claims = expected_full_token.copy()
         expected_token_with_claims.update(test_claims)
 
-        decoded = _verify_jwt(token)
+        decoded = unpack_and_verify(token)
         self.assertEqual(expected_token_with_claims, decoded)
 
     def test_malformed_token(self):
@@ -53,19 +52,8 @@ def test_malformed_token(self):
         expected_token_with_claims = expected_full_token.copy()
         expected_token_with_claims.update(test_claims)
 
-        with self.assertRaises(BadSignature):
-            _verify_jwt(token)
-
-
-def _verify_jwt(jwt_token):
-    """
-    Helper function which verifies the signature and decodes the token
-    from string back to claims form
-    """
-    keys = jwk.KEYS()
-    keys.load_jwks(settings.TOKEN_SIGNING['JWT_PUBLIC_SIGNING_JWK_SET'])
-    decoded = JWS().verify_compact(jwt_token.encode('utf-8'), keys)
-    return decoded
+        with self.assertRaises(InvalidSignatureError):
+            unpack_and_verify(token)
 
 
 @skip_unless_lms
@@ -97,33 +85,33 @@ def test_malformed_token(self):
         expected_token_with_claims = expected_full_token.copy()
         expected_token_with_claims.update(test_claims)
 
-        with self.assertRaises(BadSignature):
+        with self.assertRaises(InvalidSignatureError):
             unpack_jwt(token, test_user_id, test_now)
 
     def test_unpack_token_with_invalid_user(self):
         token = create_jwt(invalid_test_user_id, test_timeout, {}, test_now)
 
-        with self.assertRaises(Invalid):
+        with self.assertRaises(InvalidSignatureError):
             unpack_jwt(token, test_user_id, test_now)
 
     def test_unpack_expired_token(self):
         token = create_jwt(test_user_id, test_timeout, {}, test_now)
 
-        with self.assertRaises(Expired):
+        with self.assertRaises(ExpiredSignatureError):
             unpack_jwt(token, test_user_id, test_now + test_timeout + 1)
 
     def test_missing_expired_lms_user_id(self):
         payload = expected_full_token.copy()
         del payload['lms_user_id']
         token = _encode_and_sign(payload)
 
-        with self.assertRaises(MissingKey):
+        with self.assertRaises(MissingRequiredClaimError):
             unpack_jwt(token, test_user_id, test_now)
 
     def test_missing_expired_key(self):
         payload = expected_full_token.copy()
         del payload['exp']
         token = _encode_and_sign(payload)
 
-        with self.assertRaises(MissingKey):
+        with self.assertRaises(MissingRequiredClaimError):
             unpack_jwt(token, test_user_id, test_now)