add additional claims

Author: IndominusByte

fastapi_jwt_auth/auth_jwt.py

@@ -127,7 +127,8 @@ def _create_token(
         algorithm: Optional[str] = None,
         headers: Optional[Dict] = None,
         issuer: Optional[str] = None,
-        audience: Optional[Union[str,Sequence[str]]] = None
+        audience: Optional[Union[str,Sequence[str]]] = None,
+        user_claims: Optional[Dict] = {}
     ) -> str:
         """
         Create token for access_token and refresh_token (utf-8)
@@ -140,6 +141,7 @@ def _create_token(
         :param headers: valid dict for specifying additional headers in JWT header section
         :param issuer: expected issuer in the JWT
         :param audience: expected audience in the JWT
+        :param user_claims: Custom claims to include in this token. This data must be dictionary
 
         :return: Encoded token
         """
@@ -152,6 +154,8 @@ def _create_token(
             raise TypeError("audience must be a string or sequence")
         if algorithm and not isinstance(algorithm, str):
             raise TypeError("algorithm must be a string")
+        if user_claims and not isinstance(user_claims, dict):
+            raise TypeError("user_claims must be a dictionary")
 
         # Data section
         reserved_claims = {
@@ -185,7 +189,7 @@ def _create_token(
             raise
 
         return jwt.encode(
-            {**reserved_claims, **custom_claims},
+            {**reserved_claims, **custom_claims, **user_claims},
             secret_key,
             algorithm=algorithm,
             headers=headers
@@ -256,7 +260,8 @@ def create_access_token(
         algorithm: Optional[str] = None,
         headers: Optional[Dict] = None,
         expires_time: Optional[Union[timedelta,int,bool]] = None,
-        audience: Optional[Union[str,Sequence[str]]] = None
+        audience: Optional[Union[str,Sequence[str]]] = None,
+        user_claims: Optional[Dict] = {}
     ) -> str:
         """
         Create a access token with 15 minutes for expired time (default),
@@ -272,6 +277,7 @@ def create_access_token(
             algorithm=algorithm,
             headers=headers,
             audience=audience,
+            user_claims=user_claims,
             issuer=self._encode_issuer
         )
 
@@ -281,7 +287,8 @@ def create_refresh_token(
         algorithm: Optional[str] = None,
         headers: Optional[Dict] = None,
         expires_time: Optional[Union[timedelta,int,bool]] = None,
-        audience: Optional[Union[str,Sequence[str]]] = None
+        audience: Optional[Union[str,Sequence[str]]] = None,
+        user_claims: Optional[Dict] = {}
     ) -> str:
         """
         Create a refresh token with 30 days for expired time (default),
@@ -295,7 +302,8 @@ def create_refresh_token(
             exp_time=self._get_expired_time("refresh",expires_time),
             algorithm=algorithm,
             headers=headers,
-            audience=audience
+            audience=audience,
+            user_claims=user_claims
         )
 
     def _get_csrf_token(self,encoded_token: str) -> str:

tests/test_create_token.py

@@ -93,3 +93,16 @@ def test_create_token_invalid_type_data_algorithm(Authorize):
 
     with pytest.raises(TypeError,match=r"algorithm"):
         Authorize.create_refresh_token(subject=1,algorithm=1)
+
+def test_create_token_invalid_user_claims(Authorize):
+    with pytest.raises(TypeError,match=r"user_claims"):
+        Authorize.create_access_token(subject=1,user_claims="asd")
+    with pytest.raises(TypeError,match=r"user_claims"):
+        Authorize.create_refresh_token(subject=1,user_claims="asd")
+
+def test_create_valid_user_claims(Authorize):
+    access_token = Authorize.create_access_token(subject=1,user_claims={"my_access":"yeah"})
+    refresh_token = Authorize.create_refresh_token(subject=1,user_claims={"my_refresh":"hello"})
+
+    assert jwt.decode(access_token,"testing",algorithms="HS256")['my_access'] == "yeah"
+    assert jwt.decode(refresh_token,"testing",algorithms="HS256")['my_refresh'] == "hello"