Custom JWT plugin (& RelayState small fix) (#35)

* fix: check RelayState exist but not a token
- RelayState can also be a "/" (%2F char)

* feat(trigger): add custom token query
- can pas custom function to create token query in settings
- update README
- update Contributors

* Update README.md

Better CUSTOM_TOKEN_QUERY description

Co-authored-by: Mostafa Moradian <[email protected]>

* feat(jwt): better check of RelayState
- check if RelayState is a token before trying to decode it
- add test of is_jwt_well_formed

* feat(trigger): add custom jwt creator and decoder
- use custom trigger or default function for jwt management
- update README
- update tests with new functions

* fix: remove unused imports

* Update README.md

Change jwt to JWT

Co-authored-by: Mostafa Moradian <[email protected]>

Author: sgabb

AUTHORS.md

@@ -51,3 +51,4 @@ an issue.
 - [Akshit Dhar](https://github.com/akshit-wwstay)
 - [Jean Vincent](https://github.com/jean-sh)
 - [Søren Howe Gersager](https://github.com/syre)
+- [Gabrio Mauri](https://github.com/sgabb)

README.md

@@ -177,6 +177,9 @@ For IdP-initiated SSO, the user will be created if it doesn't exist, but for SP-
 | **TRIGGER.BEFORE\_LOGIN**                   | A method to be called when an existing user logs in. This method will be called before the user is logged in and after the SAML2 identity provider returns user attributes. This method should accept ONE parameter of user dict.                                                                                                                                                                                                                           | `str`            | `None`                                                                                                                                   | `my_app.models.users.before_login`                       |
 | **TRIGGER.AFTER\_LOGIN**                    | A method to be called when an existing user logs in. This method will be called after the user is logged in and after the SAML2 identity provider returns user attributes. This method should accept TWO parameters of session and user dict.                                                                                                                                                                                                               | `str`            | `None`                                                                                                                                   | `my_app.models.users.after_login`                        |
 | **TRIGGER.GET\_METADATA\_AUTO\_CONF\_URLS** | A hook function that returns a list of metadata Autoconf URLs. This can override the `METADATA_AUTO_CONF_URL` to enumerate all existing metadata autoconf URLs.                                                                                                                                                                                                                                                                                             | `str`            | `None`                                                                                                                                   | `my_app.models.users.get_metadata_autoconf_urls`         |
+| **TRIGGER.CUSTOM\_DECODE\_JWT** | A hook function to decode the user JWT. This method will be called instead of the `decode_jwt_token` default function and should return the user_model.USERNAME_FIELD. This method accepts one parameter: `token`.                                                                                                                                                                                                                                                                                             | `str`            | `None`                                                                                                                                   | `my_app.models.users.decode_custom_token`         |
+| **TRIGGER.CUSTOM\_CREATE\_JWT** | A hook function to create a custom JWT for the user. This method will be called instead of the `create_jwt_token` default function and should return the token. This method accepts one parameter: `user`.                                                                                                                                                                                                                                                                                             | `str`            | `None`                                                                                                                                   | `my_app.models.users.create_custom_token`         |
+| **TRIGGER.CUSTOM\_TOKEN\_QUERY** | A hook function to create a custom query params with the JWT for the user. This method will be called after `CUSTOM_CREATE_JWT` to populate a query and attach it to a URL; should return the query params containing the token (e.g., `?token=encoded.jwt.token`). This method accepts one parameter: `token`.                                                                                                                                                                                                                                                                                             | `str`            | `None`                                                                                                                                   | `my_app.models.users.get_custom_token_query`         |
 | **ASSERTION\_URL**                          | A URL to validate incoming SAML responses against. By default, `django-saml2-auth` will validate the SAML response's Service Provider address against the actual HTTP request's host and scheme. If this value is set, it will validate against `ASSERTION_URL` instead - perfect for when Django is running behind a reverse proxy.                                                                                                                        | `str`            | `https://example.com`                                                                                                                    |                                                          |
 | **ENTITY\_ID**                              | The optional entity ID string to be passed in the 'Issuer' element of authentication request, if required by the IDP.                                                                                                                                                                                                                                                                                                                                       | `str`            | `None`                                                                                                                                   | `https://exmaple.com/sso/acs`                            |
 | **NAME\_ID\_FORMAT**                        | Set to the string `'None'`, to exclude sending the `'Format'` property of the `'NameIDPolicy'` element in authentication requests.                                                                                                                                                                                                                                                                                                                          | `str`            | `<urn:oasis:names:tc:SAML:2.0:nameid-format:transient>`                                                                                  |                                                          |
@@ -240,6 +243,25 @@ Otherwise if you want to use your PKI key-pair to sign JWT tokens, use either of
 
 *Note:* If both PKI fields and `JWT_SECRET` are defined, the `JWT_ALGORITHM` decides which method to use for signing tokens.
 
+### Custom token triggers
+
+This is an example of the functions that could be passed to the `TRIGGER.CUSTOM_CREATE_JWT` (it uses the [DRF Simple JWT library](https://github.com/jazzband/djangorestframework-simplejwt/blob/master/docs/index.rst)) and to `TRIGGER.CUSTOM_TOKEN_QUERY`:
+
+``` python
+from rest_framework_simplejwt.tokens import RefreshToken
+
+
+def get_custom_jwt(user):
+    """Create token for user and return it"""
+    return RefreshToken.for_user(user)
+
+
+def get_custom_token_query(refresh):
+    """Create url query with refresh and access token"""
+    return "?%s%s%s%s%s" % ("refresh=", str(refresh), "&", "access=", str(refresh.access_token))
+
+```
+
 ## Customize
 
 The default permission `denied`, `error` and user `welcome` page can be overridden.

django_saml2_auth/tests/test_user.py

@@ -3,8 +3,8 @@
 import pytest
 from django.contrib.auth.models import Group
 from django_saml2_auth.exceptions import SAMLAuthError
-from django_saml2_auth.user import (create_jwt_token, create_new_user,
-                                    decode_jwt_token, get_or_create_user,
+from django_saml2_auth.user import (create_custom_or_default_jwt, create_new_user,
+                                    decode_custom_or_default_jwt, get_or_create_user,
                                     get_user, get_user_id)
 from jwt.exceptions import PyJWTError
 from pytest_django.fixtures import SettingsWrapper
@@ -306,8 +306,8 @@ def test_create_and_decode_jwt_token_success(
     """
     settings.SAML2_AUTH = saml2_settings
 
-    jwt_token = create_jwt_token("[email protected]")
-    user_id = decode_jwt_token(jwt_token)
+    jwt_token = create_custom_or_default_jwt("[email protected]")
+    user_id = decode_custom_or_default_jwt(jwt_token)
     assert user_id == "[email protected]"
 
 
@@ -347,7 +347,7 @@ def test_create_jwt_token_with_incorrect_jwt_settings(
     settings.SAML2_AUTH = saml2_settings
 
     with pytest.raises(SAMLAuthError) as exc_info:
-        create_jwt_token("[email protected]")
+        create_custom_or_default_jwt("[email protected]")
 
     assert str(exc_info.value) == error_msg
 
@@ -393,15 +393,15 @@ def test_decode_jwt_token_with_incorrect_jwt_settings(
     settings.SAML2_AUTH = saml2_settings
 
     with pytest.raises(SAMLAuthError) as exc_info:
-        decode_jwt_token("WHATEVER")
+        decode_custom_or_default_jwt("WHATEVER")
 
     assert str(exc_info.value) == error_msg
 
 
 def test_decode_jwt_token_failure():
     """Test decode_jwt_token function by passing an invalid JWT token (None, in this case)."""
     with pytest.raises(SAMLAuthError) as exc_info:
-        decode_jwt_token(None)
+        decode_custom_or_default_jwt(None)
 
     assert str(exc_info.value) == "Cannot decode JWT token."
     assert isinstance(exc_info.value.extra["exc"], PyJWTError)

django_saml2_auth/tests/test_utils.py

@@ -2,7 +2,7 @@
 from django.http import HttpRequest, HttpResponse
 from django.urls import NoReverseMatch
 from django_saml2_auth.exceptions import SAMLAuthError
-from django_saml2_auth.utils import exception_handler, get_reverse, run_hook
+from django_saml2_auth.utils import exception_handler, get_reverse, run_hook, is_jwt_well_formed
 
 
 def divide(a: int, b: int = 1) -> int:
@@ -124,3 +124,12 @@ def test_exception_handler_handle_exception():
     contents = result.content.decode("utf-8")
     assert result.status_code == 500
     assert "Reason: Internal world error!" in contents
+
+
+def test_jwt_well_formed():
+    """Test if passed RelayState is a well formed JWT"""
+    token = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiI0MjQyIiwibmFtZSI6Ikplc3NpY2EgVGVtcG9yYWwiLCJuaWNrbmFtZSI6Ikplc3MifQ.EDkUUxaM439gWLsQ8a8mJWIvQtgZe0et3O3z4Fd_J8o'  # noqa
+    res = is_jwt_well_formed(token)  # True
+    assert res is True
+    res = is_jwt_well_formed('/')  # False
+    assert res is False

django_saml2_auth/user.py

@@ -327,6 +327,44 @@ def create_jwt_token(user_id: str) -> Optional[str]:
     return jwt.encode(payload, secret, algorithm=jwt_algorithm)
 
 
+def create_custom_or_default_jwt(user: Union[str, Type[Model]]):
+    """Create a new JWT token, eventually using custom trigger
+
+    Args:
+        user (dict or str): User instance or User's username or email based on User.USERNAME_FIELD
+
+    Returns:
+        Optional[str]: JWT token
+    """
+    saml2_auth_settings = settings.SAML2_AUTH
+    custom_create_jwt_trigger = dictor(saml2_auth_settings, "TRIGGER.CUSTOM_CREATE_JWT")
+
+    # If user is the id (user_model.USERNAME_FIELD), set it as user_id
+    user_id = None
+    if isinstance(user, str):
+        user_id = user
+
+    # Check if there is a custom trigger for creating the JWT and URL query
+    if custom_create_jwt_trigger:
+        target_user = user
+        # If user is user_id, get user instance
+        if user_id:
+            user_model = get_user_model()
+            user = {
+                user_model.USERNAME_FIELD: user_id
+            }
+            target_user = get_user(user)
+        jwt_token = run_hook(custom_create_jwt_trigger, target_user)
+    else:
+        # If user_id is not set, retrieve it from user instance
+        if not user_id:
+            user_id = getattr(user, user_model.USERNAME_FIELD)
+        # Create a new JWT token with PyJWT
+        jwt_token = create_jwt_token(user_id)
+
+    return jwt_token
+
+
 def decode_jwt_token(jwt_token: str) -> Optional[str]:
     """Decode a JWT token
 
@@ -366,3 +404,24 @@ def decode_jwt_token(jwt_token: str) -> Optional[str]:
             "reason": "Cannot decode JWT token.",
             "status_code": 500
         })
+
+
+def decode_custom_or_default_jwt(jwt_token: str) -> Optional[str]:
+    """Decode a JWT token, eventually using custom trigger
+
+    Args:
+        jwt_token (str): The token to decode
+
+    Raises:
+        SAMLAuthError: Cannot decode JWT token.
+
+    Returns:
+        Optional[str]: A user_id as str or None.
+    """
+    saml2_auth_settings = settings.SAML2_AUTH
+    custom_decode_jwt_trigger = dictor(saml2_auth_settings, "TRIGGER.CUSTOM_DECODE_JWT")
+    if custom_decode_jwt_trigger:
+        user_id = run_hook(custom_decode_jwt_trigger, jwt_token)
+    else:
+        user_id = decode_jwt_token(jwt_token)
+    return user_id

django_saml2_auth/utils.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+import base64
 from functools import wraps
 from typing import Any, Callable, Iterable, Mapping, Optional, Tuple, Union
 
@@ -157,3 +158,29 @@ def wrapper(request: HttpRequest) -> HttpResponse:
             result = handle_exception(exc, request)
         return result
     return wrapper
+
+
+def is_jwt_well_formed(jwt: str):
+    """Check if JWT is well formed
+
+    Args:
+        jwt (str): Json Web Token
+
+    Returns:
+        Boolean: True if JWT is well formed, otherwise False
+    """
+    if isinstance(jwt, str):
+        # JWT should contain three segments, separated by two period ('.') characters.
+        jwt_segments = jwt.split('.')
+        if len(jwt_segments) == 3:
+            jose_header = jwt_segments[0]
+            # base64-encoded string length should be a multiple of 4
+            if len(jose_header) % 4 == 0:
+                try:
+                    jh_decoded = base64.b64decode(jose_header).decode('utf-8')
+                    if jh_decoded and jh_decoded.find('JWT') > -1:
+                        return True
+                except Exception:
+                    return False
+    # If tests not passed return False
+    return False

django_saml2_auth/views.py

@@ -5,7 +5,6 @@
 
 import urllib.parse as urlparse
 from urllib.parse import unquote
-import json
 
 from dictor import dictor
 from django import get_version
@@ -26,8 +25,8 @@
                                     extract_user_identity, get_assertion_url,
                                     get_default_next_url, get_saml_client)
 from django_saml2_auth.user import (
-    get_or_create_user, create_jwt_token, decode_jwt_token, get_user_id)
-from django_saml2_auth.utils import exception_handler, get_reverse, run_hook
+    create_custom_or_default_jwt, decode_custom_or_default_jwt, get_or_create_user, get_user_id)
+from django_saml2_auth.utils import exception_handler, get_reverse, is_jwt_well_formed, run_hook
 from pkg_resources import parse_version
 
 
@@ -71,11 +70,15 @@ def acs(request: HttpRequest):
 
     next_url = request.session.get("login_next_url") or get_default_next_url()
 
-    # If RelayState params is passed, it is a JWT token that identifies the user trying to login
-    # via sp_initiated_login endpoint
+    # A RelayState is an HTTP parameter that can be included as part of the SAML request and SAML response;
+    # usually is meant to be an opaque identifier that is passed back without any modification or inspection,
+    # and it is used to specify additional information to the SP or the IdP.
+    # If RelayState params is passed, it could be JWT token that identifies the user trying to login
+    # via sp_initiated_login endpoint, or it could be a URL used for redirection.
     relay_state = request.POST.get("RelayState")
-    if relay_state:
-        redirected_user_id = decode_jwt_token(relay_state)
+    relay_state_is_token = is_jwt_well_formed(relay_state)
+    if relay_state_is_token:
+        redirected_user_id = decode_custom_or_default_jwt(relay_state)
 
         # This prevents users from entering an email on the SP, but use a different email on IdP
         if get_user_id(user) != redirected_user_id:
@@ -97,10 +100,14 @@ def acs(request: HttpRequest):
     use_jwt = dictor(saml2_auth_settings, "USE_JWT", False)
     if use_jwt and target_user.is_active:
         # Create a new JWT token for IdP-initiated login (acs)
-        jwt_token = create_jwt_token(target_user.email)
-        # Use JWT auth to send token to frontend
-        query = f"?token={jwt_token}"
+        jwt_token = create_custom_or_default_jwt(target_user)
+        custom_token_query_trigger = dictor(saml2_auth_settings, "TRIGGER.CUSTOM_TOKEN_QUERY")
+        if custom_token_query_trigger:
+            query = run_hook(custom_token_query_trigger, jwt_token)
+        else:
+            query = f"?token={jwt_token}"
 
+        # Use JWT auth to send token to frontend
         frontend_url = dictor(saml2_auth_settings, "FRONTEND_URL", next_url)
 
         return HttpResponseRedirect(frontend_url + query)
@@ -134,9 +141,9 @@ def sp_initiated_login(request: HttpRequest) -> HttpResponseRedirect:
     # User must be created first by the IdP-initiated SSO (acs)
     if request.method == "GET":
         if request.GET.get("token"):
-            user_id = decode_jwt_token(request.GET.get("token"))
+            user_id = decode_custom_or_default_jwt(request.GET.get("token"))
             saml_client = get_saml_client(get_assertion_url(request), acs, user_id)
-            jwt_token = create_jwt_token(user_id)
+            jwt_token = create_custom_or_default_jwt(user_id)
             _, info = saml_client.prepare_for_authenticate(sign=False, relay_state=jwt_token)
             redirect_url = dict(info["headers"]).get("Location", "")
             if not redirect_url: