From 2a4f84a36324e1286f558a693f5b7416c2e8cb75 Mon Sep 17 00:00:00 2001
From: Logan Deniston <logan.deniston@gmail.com>
Date: Tue, 27 Jun 2023 00:47:00 -0400
Subject: [PATCH] Added ability to use datetime objects that are timezone aware
 for time based claims.

---
 jwcrypto/jwt.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/jwcrypto/jwt.py b/jwcrypto/jwt.py
index b7a29b8..f27c229 100644
--- a/jwcrypto/jwt.py
+++ b/jwcrypto/jwt.py
@@ -3,6 +3,7 @@
 import copy
 import time
 import uuid
+from datetime import datetime
 
 from deprecated import deprecated
 
@@ -172,7 +173,9 @@ def __init__(self, header=None, claims=None, jwt=None, key=None,
         """Creates a JWT object.
 
         :param header: A dict or a JSON string with the JWT Header data.
-        :param claims: A dict or a string with the JWT Claims data.
+        :param claims: A dict or a string with the JWT Claims data. If
+         the 'exp' or 'nbf' are datetime objects, they are automatically
+         converted to integer unix timestamps. Otherwise, they are left as is.
         :param jwt: a 'raw' JWT token
         :param key: A (:class:`jwcrypto.jwk.JWK`) key to deserialize
          the token. A (:class:`jwcrypto.jwk.JWKSet`) can also be used.
@@ -225,6 +228,14 @@ def __init__(self, header=None, claims=None, jwt=None, key=None,
             self._check_claims = check_claims
 
         if claims is not None:
+            # Check for datetime objects
+            if 'exp' in claims and isinstance(claims['exp'], datetime):
+                # Check if timezone aware
+                claims['exp'] = self._check_and_convert_dt(claims['exp'], 'exp')
+            if 'nbf' in claims and isinstance(claims['nbf'], datetime):
+                # Check if timezone aware
+                claims['nbf'] = self._check_and_convert_dt(claims['nbf'], 'nbf')
+
             self.claims = claims
 
         if jwt is not None:
@@ -388,6 +399,18 @@ def expected_type(self, v):
         else:
             raise ValueError("Invalid value, must be 'JWS' or 'JWE'")
 
+    def _check_and_convert_dt(self, dt, claim_prop):
+        if (
+            dt.tzinfo is not None
+            and dt.tzinfo.utcoffset(dt) is not None
+        ):
+            dt_timestamp = int(dt.timestamp())
+            return dt_timestamp
+        else:
+            raise ValueError(
+                f"'{claim_prop}' datetime object must be timezone aware"
+            )
+
     def _add_optional_claim(self, name, claims):
         if name in claims:
             return