Customizing mozilla-django-oidc
Tags: python, django, okta, tutorial
Published: 2022-12-30 12:07:27

### Prerequisite:
> In order to implement OIDC authentication using Django and mozilla-django-oidc, I would recommend reading this article: [**How to implement OIDC authentication with Django and Okta**]. {% embed https://dev.to/hesbon/oidc-oauth2-authentication-using-django-and-mozilla-django-oidc-with-okta-4jll %}

- _The [mozilla-django-oidc](https://mozilla-django-oidc.readthedocs.io/en/stable/) package provides a Django authentication backend that allows you to use OIDC as the authentication mechanism for your Django application. It handles the details of verifying the user's identity and obtaining the necessary information from the OIDC provider. The backend can be used to authenticate users in Django's built-in authentication system, or you can use it to add OIDC support to your own custom authentication backend._
- _Ensure you have an Okta developer account and configure an application. Follow the steps detailed in the article above._

> Note: _Here's a link to the final project on Github._ {% embed https://github.com/Hesbon5600/oidc-connect %}
- In the Okta application you created, ensure you have enabled `token refresh` as shown below. ![Enable token rfresh](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ek59vn3htzmjew4ygh5c.png)

---

### Introduction

In this article, I am going to address three main areas:
1. Storing the user's OIDC `access_token`, `id_token`, and `refresh_token`. By default the package only allows us to store the access token and id token. Since we want to implement session refresh and logout, we will need to store the "access" and "refresh" tokens.
2. OIDC Logout. `mozilla_django_oidc` logout only terminates the existing session in our Django app. We want to also terminate the session in the Okta authorization server. 
3. Implementing our own session refresh middleware based on the `mozilla_django_oidc.middleware.SessionRefresh` middleware. By default, the package forces a user to login when their token expires. We might want to automatically refresh the user's `access token` as long as the `refresh token` is still valid. We will also rotate the `refresh token` after each use.

> When using the Okta authorization server, the lifetime of the JWT tokens is hard-coded to the following values: `ID Token`: 60 minutes. `Access Token`: 60 minutes. `Refresh Token`: 100 days.

---

#### Part 1: Storing access and refresh tokens

- For this, we need to create our own custom authentication backend that subclasses `mozilla_django_oidc.auth.OIDCAuthenticationBackend` and override the `authenticate` method.

```python
# oidc_app/core/backends.py

import logging

from django.core.exceptions import SuspiciousOperation
from django.urls import reverse
from mozilla_django_oidc.auth import OIDCAuthenticationBackend
from mozilla_django_oidc.utils import absolutify


LOGGER = logging.getLogger(__name__)


class CustomOIDCAuthenticationBackend(OIDCAuthenticationBackend):
    """Custom OIDC authentication backend."""

    def authenticate(self, request, **kwargs):
        """Authenticates a user based on the OIDC code flow."""

        self.request = request
        if not self.request:
            return None

        state = self.request.GET.get("state")
        code = self.request.GET.get("code")
        nonce = kwargs.pop("nonce", None)

        if not code or not state:
            return None

        reverse_url = self.get_settings(
            "OIDC_AUTHENTICATION_CALLBACK_URL", "oidc_authentication_callback"
        )

        token_payload = {
            "client_id": self.OIDC_RP_CLIENT_ID,
            "client_secret": self.OIDC_RP_CLIENT_SECRET,
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": absolutify(self.request, reverse(reverse_url)),
        }

        # Get the token
        token_info = self.get_token(token_payload)
        id_token = token_info.get("id_token")
        access_token = token_info.get("access_token")
        refresh_token = token_info.get("refresh_token")

        # Validate the token
        payload = self.verify_token(id_token, nonce=nonce)

        if payload:
            self.store_tokens(access_token, id_token, refresh_token) # <--- HERE: store tokens
            try:
                return self.get_or_create_user(access_token, id_token, payload)
            except SuspiciousOperation as exc:
                LOGGER.warning("failed to get or create user: %s", exc)
                return None

        return None

    def store_tokens(self, access_token, id_token, refresh_token): # <--- HERE: store tokens
        """Store OIDC tokens."""
        session = self.request.session

        if self.get_settings("OIDC_STORE_ACCESS_TOKEN", True):
            session["oidc_access_token"] = access_token

        if self.get_settings("OIDC_STORE_ID_TOKEN", False):
            session["oidc_id_token"] = id_token

        if self.get_settings("OIDC_STORE_REFRESH_TOKEN", True): # <--- HERE: Add refresh token option
            session["oidc_refresh_token"] = refresh_token
```
- Update the settings file to add the option of storing the tokens and add the new authentication backend. We also need to explicitly tell "mozilla_django_oidc" to set the token expiry time to `1 hour` from the authentication time. Furthermore, since Okta does not automatically issue a refresh token, you need to add `offline_access` scope to your `OIDC_RP_SCOPES` value as follows:

```python
# oidc_app/settings.py

...
AUTHENTICATION_BACKENDS = (
    "oidc_app.core.backends.CustomOIDCAuthenticationBackend",
    "django.contrib.auth.backends.ModelBackend",
)
...
OIDC_RENEW_ID_TOKEN_EXPIRY_SECONDS = 60 * 60  # 1 hour
OIDC_STORE_ACCESS_TOKEN = os.environ.get("OIDC_STORE_ACCESS_TOKEN", True) # Store the access token in the OIDC backend
OIDC_STORE_ID_TOKEN = os.environ.get("OIDC_STORE_ID_TOKEN", True) # Store the ID token in the OIDC backend
OIDC_STORE_REFRESH_TOKEN = os.environ.get("OIDC_STORE_REFRESH_TOKEN", True) # Store the refresh token in the OIDC backend
OIDC_RP_SCOPES = os.environ.get("OIDC_RP_SCOPES", "openid profile email offline_access")  # The OIDC scopes to request

...
```
> That's it. The token sill be stored in the user session whenever the user authenticates.

---

#### Part 2: OIDC Logout

> Although both refresh tokens and access tokens have an expiration time, it is highly advised to revoke these tokens once they aren’t needed (e.g. in case the user logs out of your application).

- Create a `logout view` that handles; django logout and oidc  logout.
```python
# oidc_app/authentication/views.py
...
from django.contrib.auth import logout
from django.http import HttpResponseRedirect
from django.contrib.auth.views import LogoutView
from django.utils.decorators import method_decorator
from django.views.decorators.cache import never_cache
from oidc_app.core.logout import oidc_logout # to be implemented
...
...
class LogoutViewSet(LogoutView):
    """
    API endpoint that handles user logout
    """
    @method_decorator(never_cache)
    def dispatch(self, request, *args, **kwargs):
        """
        Override the dispatch method to add the okta logout functionality
        """

        oidc_logout(request) # This is the function that does the oidc logout
        logout(request) # This is the django logout
        redirect_to = self.get_success_url()
        if redirect_to != request.get_full_path():
            # Redirect to target page once the session has been cleared.
            return HttpResponseRedirect(redirect_to)
        return super().dispatch(request, *args, **kwargs)
```

- Update the `urls` file to include the new logout viewset
```python
# oidc_app/urls.py
...
from oidc_app.authentication import views as auth_views

urlpatterns = [
   ...
    path("logout", auth_views.LogoutViewSet.as_view(), name="logout"),
```

**Okta Logout**
- Add the `token revoke` endpoint to the settings file as follows.

```python
# oidc_app/settings.py
OIDC_OP_TOKEN_REVOKE_ENDPOINT = (
    f"https://{OKTA_DOMAIN}/oauth2/default/v1/revoke"  # The OIDC token revocation endpoint
)
```
- With all the configurations in place, we need to make a request to Okta and revoking the access and refresh tokens

```python
# oidc_app/core/oidc_logout.py

import requests
import logging
from django.conf import settings

LOGGER = logging.getLogger(__name__)


def revoke_token(token_type, token):
    """Revoke an OIDC token."""

    token_revoke_payload = {
        "client_id": settings.OIDC_RP_CLIENT_ID,
        "client_secret": settings.OIDC_RP_CLIENT_SECRET,
        "token": token,
        "token_type_hint": token_type,
    }

    try:
        response = requests.post(
            settings.OIDC_OP_TOKEN_REVOKE_ENDPOINT, data=token_revoke_payload
        )
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        LOGGER.error("Failed to revoke token: %s", e)


def oidc_logout(request):
    """Logout the user."""

    token_types = {
        "refresh_token": request.session.get("oidc_refresh_token"),
        "access_token": request.session.get("oidc_access_token"),
    }

    for token_type, token in token_types.items():
        if token:
            LOGGER.info("Revoking token of type %s", token_type)
            revoke_token(token_type, token)
```

> That's it!! When your user logs out, the access and refresh tokens will be revoked by Okta.


---

#### Part 3: Session Refresh middleware.

> Access and ID tokens are JSON web tokens that are valid for a specific number of seconds. Typically, a user needs a new access token when they attempt to access a resource for the first time or after the previous access token that was granted to them expires. A refresh token is a special token that is used to obtain additional access tokens. This allows you to have short-lived access tokens without having to collect credentials every time one expires. You request a refresh token alongside the access and/or ID tokens as part of a user's initial authentication and authorization flow. Applications must then securely store refresh tokens since they allow users to remain authenticated.

- Ideally, you would want to check the user session validity status whenever they make a request. Hence, we are going to subclass `mozilla_django_oidc.middleware.SessionRefresh` and handle the session refresh if the `access token` has expired.

```python
# oidc_app/core/middleware.py

import logging
import time
from re import Pattern as re_Pattern

import requests
from django.urls import reverse
from django.utils.functional import cached_property
from mozilla_django_oidc.middleware import SessionRefresh as OIDCSessionRefresh

LOGGER = logging.getLogger(__name__)


class OIDCSessionRefreshMiddleware(OIDCSessionRefresh):
    @cached_property
    def exempt_urls(self):
        """Generate and return a set of url paths to exempt from SessionRefresh

        This takes the value of ``settings.OIDC_EXEMPT_URLS`` and appends three
        urls that mozilla-django-oidc uses. These values can be view names or
        absolute url paths.

        :returns: list of url paths (for example "/oidc/callback/")

        """
        exempt_urls = []
        for url in self.OIDC_EXEMPT_URLS:
            if not isinstance(url, re_Pattern):
                exempt_urls.append(url)
        return set(
            [url if url.startswith("/") else reverse(url) for url in exempt_urls]
        )

    def refresh_session(self, request):
        """Refresh the session with new data from the request session store."""

        refresh_token = request.session.get("oidc_refresh_token", None)

        token_refresh_payload = {
            "refresh_token": refresh_token,
            "client_id": self.get_settings("OIDC_RP_CLIENT_ID"),
            "client_secret": self.get_settings("OIDC_RP_CLIENT_SECRET"),
            "grant_type": "refresh_token",
        }

        try:
            response = requests.post(
                self.get_settings("OIDC_OP_TOKEN_ENDPOINT"), data=token_refresh_payload
            )
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            LOGGER.error("Failed to refresh session: %s", e)
            return False
        data = response.json()
        request.session.update(
            {
                "oidc_access_token": data.get("access_token"),
                "oidc_id_token_expiration": time.time() + data.get("expires_in"),
                "oidc_refresh_token": data.get("refresh_token"),
            }
        )
        return True

    def process_request(self, request):
        if not self.is_refreshable_url(request):
            LOGGER.debug("request is not refreshable")
            return

        expiration = request.session.get("oidc_id_token_expiration", 0)
        now = time.time()
        if expiration > now:
            # The id_token is still valid, so we don't have to do anything.
            LOGGER.debug("id token is still valid (%s > %s)", expiration, now)
            return

        LOGGER.debug("id token has expired")
        if not self.refresh_session(request):
            # If we can't refresh the session, then we need to reauthenticate the user.
            # As per the default OIDCSessionRefresh implementation.
            return super().process_request(request)

        LOGGER.debug("session refreshed")
```

- We need to update the `MIDDLEWARE` settings to include the new middleware we just created. We also have a couple of URLs that needs to be skipped when doing the session refresh: `oidc_authentication_callback`, `oidc_authentication_callback`, and `logout`. Update your settings file with the list of these urls.
 
```python
# oidc_app/settings.py

MIDDLEWARE = [
    ...
   "django.contrib.auth.middleware.AuthenticationMiddleware",
    "oidc_app.core.middleware.OIDCSessionRefreshMiddleware",
    ...
]

OIDC_EXEMPT_URLS = [
    "oidc_authentication_init",
    "oidc_authentication_callback",
    "logout",
]
```

> **Thats it!** You now have a custom backend that stores the refresh and access token, a middleware that handles fetching a new refresh and access token as well as a logout view that invalidates the access and refresh tokens.

**Next Up:** Implement the Authorization Code with PKCE flow in Okta. Stay tunned.

Feel free to leave a comment or suggestion. Thank you!