From 1580f993fe26fb74b95249d05f5bd482ff6f1900 Mon Sep 17 00:00:00 2001
From: Sara Ghaemi <sara.ghaemi@telus.com>
Date: Tue, 7 May 2024 11:22:17 -0400
Subject: [PATCH] Updated JWT handler to support PEM public key

---
 litellm/proxy/auth/handle_jwt.py | 34 ++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index 1324c2c594..c12f48cc11 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -15,6 +15,9 @@ from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_UserTable
 from litellm.proxy.utils import PrismaClient
 from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
 from typing import Optional
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 
 
 class JWTHandler:
@@ -142,8 +145,8 @@ class JWTHandler:
                 public_key = keys[0]
         elif len(keys) > 1:
             for key in keys:
-                if kid is not None and key["kid"] == kid:
-                    public_key = key
+                if kid is not None and key == kid:
+                    public_key = keys[key]
 
         if public_key is None:
             raise Exception(
@@ -191,6 +194,33 @@ class JWTHandler:
                 raise Exception("Token Expired")
             except Exception as e:
                 raise Exception(f"Validation fails: {str(e)}")
+        elif public_key is not None and isinstance(public_key, str):
+            audience = os.getenv("JWT_AUDIENCE")
+            if audience is None:
+                raise Exception("Missing JWT Audience from environment.")
+            try:
+                cert = x509.load_pem_x509_certificate(public_key.encode(), default_backend())
+
+                # Extract public key
+                key = cert.public_key().public_bytes(
+                    serialization.Encoding.PEM,
+                    serialization.PublicFormat.SubjectPublicKeyInfo
+                )
+
+                # decode the token using the public key
+                payload = jwt.decode(
+                    token,
+                    key,
+                    algorithms=["RS256"],
+                    audience=audience,
+                )
+                return payload
+
+            except jwt.ExpiredSignatureError:
+                # the token is expired, do something to refresh it
+                raise Exception("Token Expired")
+            except Exception as e:
+                raise Exception(f"Validation fails: {str(e)}")
 
         raise Exception("Invalid JWT Submitted")