From e481e79ff44892a887c38dbdbf2aa99f48f2394e Mon Sep 17 00:00:00 2001
From: jamesstottmoj <james.stott@digital.justice.gov.uk>
Date: Fri, 20 Dec 2024 10:29:39 +0000
Subject: [PATCH] Replaced python-jose with PyJWT

---
 controlpanel/api/jwt_auth.py |  6 +++---
 controlpanel/jwt.py          | 41 +++++++++++-------------------------
 requirements.txt             |  2 +-
 3 files changed, 16 insertions(+), 33 deletions(-)

diff --git a/controlpanel/api/jwt_auth.py b/controlpanel/api/jwt_auth.py
index 4f5676845..031a887d0 100644
--- a/controlpanel/api/jwt_auth.py
+++ b/controlpanel/api/jwt_auth.py
@@ -7,7 +7,7 @@
 
 # First-party/Local
 from controlpanel.api.models import User
-from controlpanel.jwt import JWT, JWTDecodeError
+from controlpanel.jwt import JWT, DecodeError
 
 M2M_CLAIM_FLAG = "client-credentials"
 
@@ -69,7 +69,7 @@ def authenticate(self, request):
         else:
             try:
                 jwt.validate()
-            except JWTDecodeError:
+            except DecodeError:
                 return None
 
         return self._get_client(jwt), None
@@ -89,7 +89,7 @@ def _get_client(self, jwt):
                 return AuthenticatedServiceClient(jwt.payload)
             else:
                 raise exceptions.AuthenticationFailed()
-        except JWTDecodeError:
+        except DecodeError:
             raise exceptions.AuthenticationFailed(
                 "Failed to be authenticated due to JWT decoder error!"
             )
diff --git a/controlpanel/jwt.py b/controlpanel/jwt.py
index 6e47c1292..04cdb301f 100644
--- a/controlpanel/jwt.py
+++ b/controlpanel/jwt.py
@@ -1,19 +1,13 @@
 # Third-party
-import requests
+import jwt
 import structlog
 from django.conf import settings
-from jose import jwt
-from jose.exceptions import JWTError
-from requests.exceptions import RequestException
+from jwt.exceptions import DecodeError, InvalidTokenError, PyJWKClientError
 from rest_framework import HTTP_HEADER_ENCODING
 
 log = structlog.getLogger(__name__)
 
 
-class JWTDecodeError(Exception):
-    pass
-
-
 class JWT:
     def __init__(self, raw_token):
         self._header = None
@@ -25,7 +19,7 @@ def __init__(self, raw_token):
             "algorithms": [settings.OIDC_RP_SIGN_ALGO],
             "audience": settings.OIDC_CPANEL_API_AUDIENCE,
             "options": {
-                "require_sub": True,
+                "require": ["sub"],
             },
         }
 
@@ -37,7 +31,7 @@ def header(self):
         if not self._header:
             try:
                 self._header = jwt.get_unverified_header(self._raw_token)
-            except jwt.JWTError:
+            except (DecodeError, InvalidTokenError):
                 return None
         return self._header
 
@@ -45,22 +39,11 @@ def header(self):
     def jwk(self):
         if not self._jwk and self.header:
             try:
-                response = requests.get(self.jwks_url, verify=False)
-                response.raise_for_status()
-            except RequestException as error:
-                raise JWTDecodeError(f"Failed fetching JWK: {error}")
-
-            jwks = response.json()
+                jwks_client = jwt.PyJWKClient(self.jwks_url)
+                self._jwk = jwks_client.get_signing_key_from_jwt(self._raw_token).key
 
-            for jwk in jwks.get("keys", []):
-                if jwk["kid"] == self.header["kid"]:
-                    self._jwk = jwk
-                    return self._jwk
-
-            raise JWTDecodeError(
-                f'No JWK with id {self.header["kid"]} found at {self.jwks_url} '
-                f"while decoding {self._raw_token}"
-            )
+            except PyJWKClientError as error:
+                raise DecodeError(f"Failed fetching JWK: {error}")
 
         return self._jwk
 
@@ -73,8 +56,8 @@ def payload(self):
                     key=self.jwk,
                     **self.decode_options,
                 )
-            except (JWTError, KeyError) as error:
-                raise JWTDecodeError(f"Failed decoding JWT: {error}")
+            except (DecodeError, KeyError) as error:
+                raise DecodeError(f"Failed decoding JWT: {error}")
         return self._payload
 
     def validate(self):
@@ -84,8 +67,8 @@ def validate(self):
                 key=self.jwk,
                 **self.decode_options,
             )
-        except (JWTError, KeyError) as error:
-            raise JWTDecodeError(f"Failed decoding JWT: {error}")
+        except (DecodeError, KeyError) as error:
+            raise DecodeError(f"Failed decoding JWT: {error}")
 
     @classmethod
     def from_auth_header(cls, request):
diff --git a/requirements.txt b/requirements.txt
index c7ad14afb..d46d0d38f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -27,11 +27,11 @@ model-bakery==1.17.0
 moto[all]==5.0.18
 mozilla-django-oidc==4.0.1
 psycopg2-binary==2.9.9
+PyJWT==2.10.1
 PyNaCl==1.5.0
 pytest==8.0.0
 pytest-django==4.9.0
 python-dotenv==1.0.1
-python-jose==3.3.0
 pyyaml==6.0.2
 rules==3.3
 sentry-sdk==2.19.2