Improved django-revproxy ProxyView: Handle root URLs in content, Location and clean django's sessionid cookie

pimped_proxy_view.py

import re
from urllib.parse import urlparse

from revproxy.utils import is_html_content_type
from revproxy.views import ProxyView

from django.conf import settings
from django.http import SimpleCookie
from django.urls import resolve, reverse


class GetBasePathMixin:
    def get_base_path(self, request):
        """
        Return the view django's path without `path` kwargs
        /django/proxy/42/root/upstream/path
        => /django/proxy/42/root/
        """
        resolver_match = resolve(request.path)
        kwargs = {**resolver_match.kwargs, "path": ""}
        return reverse(resolver_match.url_name, kwargs=kwargs)


class HandleRootUrlMixin(GetBasePathMixin):
    """
    This mixin  manage case when HTML content is referring folder outside given upstream path.
    ex:
    self.upstream = "http://domain.xyz/dir1/"
    <img src="/images/image.jpg" />
    It will ensure the reverse proxy will fetch http://domain.xyz/images/image.jpg
    instead of http://domain.xyz/dir1/images/image.jpg.
    """

    replace_root_url = True
    root_url_marker = "{ROOT}"
    root_url_html_elements_regex = r"((?:src|href|action) *= *[\"']?)"

    def replace_absolute_urls(self, content: str, root_path: str):
        content = re.sub(self.root_url_html_elements_regex + "/", rf"\g<1>{root_path}", content)
        root_url = self.get_root_url()
        return re.sub(self.root_url_html_elements_regex + root_url, rf"\g<1>{root_path}", content)

    def get_root_marked_path(self, request):
        return self.get_base_path(request) + f"{self.root_url_marker}/"

    def get_root_url(self):
        parsed_result = urlparse(self.upstream)
        return f"{parsed_result.scheme}://{parsed_result.netloc}/"

    def dispatch(self, request, path, *args, **kwargs):
        if not self.replace_root_url:
            return super().dispatch(request, path, *args, **kwargs)

        if self.root_url_marker in path:
            self.upstream = self.get_root_url()
            path = path.replace(f"{self.root_url_marker}/", "")

        response = super().dispatch(request, path, *args, **kwargs)

        content_type = response.get("Content-Type")
        if response.streaming or not is_html_content_type(content_type):
            return response

        root_path = self.get_root_marked_path(request)
        response.content = self.replace_absolute_urls(response.content.decode(response.charset), root_path).encode()
        del response["Content-Length"]

        return response


class HandleRootRedirect(GetBasePathMixin):
    """
    If response has Location header which redirect to folder outside upstream path,
    It force the Location header to be the revproxy view base path.
    Ex:
    Location: /new-path
    => Location: /django/proxy/42/root/new-path
    """

    handle_root_redirect = True

    def _replace_host_on_redirect_location(self, request, proxy_response):
        super()._replace_host_on_redirect_location(request, proxy_response)
        location = proxy_response.headers.get("Location")
        if not location or not location.startswith("/"):
            return
        location = self.get_base_path(request) + location[1:]
        proxy_response.headers["Location"] = location


class NoCookieMixin:
    """
    Remove cookie transmission between client and upstream
    """

    no_cookie = True

    def get_request_headers(self, *args, **kwargs):
        request_headers = super().get_request_headers()
        if not self.no_cookie:
            return request_headers

        if "Cookie" in request_headers:
            del request_headers["Cookie"]
        return request_headers

    def dispatch(self, *args, **kwargs):
        response = super().dispatch(*args, **kwargs)
        if self.no_cookie:
            response.cookies.clear()
        return response


class HideSessionIdCookieMixin:
    """
    Remove django's sessionid's cookie transmission between client and upstream
    """

    hide_session_id_cookie = True

    def get_session_id_cookie_key(self):
        return settings.SESSION_COOKIE_NAME

    def get_request_headers(self, *args, **kwargs):
        request_headers = super().get_request_headers()
        if not self.hide_session_id_cookie:
            return request_headers

        if "Cookie" in request_headers:
            cookie = SimpleCookie(request_headers["Cookie"])
            session_id_key = self.get_session_id_cookie_key()
            if session_id_key in cookie:
                del cookie[session_id_key]
                request_headers["Cookie"] = cookie.output(header="", sep=";")

        return request_headers

    def dispatch(self, *args, **kwargs):
        response = super().dispatch(*args, **kwargs)
        if self.hide_session_id_cookie:
            session_id_key = self.get_session_id_cookie_key()
            if session_id_key in response.cookies:
                del response.cookies[session_id_key]
        return response


class PimpedProxyView(HandleRootRedirect, HandleRootUrlMixin, HideSessionIdCookieMixin, NoCookieMixin, ProxyView):
    """
    Override of ProxyView with features like :
    - Handle URL referring upstream root
    - Cleaning cookies or django's sessionid cookie

    Features are disabled by default, enable with :

    replace_root_url=True
    handle_root_redirect=True
    no_cookie=True
    hide_session_id_cookie=True
    """

    replace_root_url = False
    handle_root_redirect = False
    no_cookie = False
    hide_session_id_cookie = False