Decorators depending on request method types about django-csp HOT 3 OPEN

from functools import wraps
from django.http import HttpResponse
from django.http import StreamingHttpResponse
from django_csp.middleware import CSPMiddleware


def csp_exclude(methods=[]):
    def decorator(func):
        @wraps(func)
        def inner_func(request, args, **kwargs):
            if request.method in methods:
                response = func(request,args, kwargs)
                return response
            else:
                # Exclude CSP headers for this request method
                response = HttpResponse()
                setattr(response, '_csp_exempt', True)
                return response

        return inner_func

    return decorator


def csp_update(methods=[], kwargs):
    def decorator(func):
        @wraps(func)
        def inner_func(request, *args, **kwargs):
            if request.method in methods:

                middleware = CSPMiddleware(get_response=None)
                policy = middleware.build_policy(base_policy={})
                policy.update(kwargs)
                # Set the updated policy on the response object
                response = func(request, args, **kwargs)
                response._csp_update = policy
                return response
            else:
                response = func(request,args, **kwargs)
                return response

        return inner_func

    return decorator

# Usage Example
@csp_update(methods=['POST'], kwargs={'script-src': ["'self'", "https://example.com/)
def my_view(request):
    if request.method == 'POST':
        # perform some actions/operations...
        return HttpResponse("POST request successful!")
    else:
        return HttpResponse("GET request successful!")

from django.test import RequestFactory, TestCase
from django.http import HttpResponse
from django_csp.middleware import CSPMiddleware
from myapp.views import my_view, csp_update

class CSPMiddlewareTest(TestCase):
    def setUp(self):
        self.factory = RequestFactory()

    def test_csp_update(self):
        # Create a request object
        request = self.factory.post('/my-view/')
        kwargs = {'script-src': ["'self'", "https://example.com/%22]%7D

        # Create a response object with the CSP header
        response = HttpResponse()
        middleware = CSPMiddleware(get_response=None)
        policy = middleware.build_policy(base_policy={})
        policy.update(kwargs)
        response['Content-Security-Policy'] = middleware._build_policy(policy)

        # Call the view decorated with @csp_update and check that the returned response object has the updated CSP header
        view = csp_update(methods=['POST'], kwargs=kwargs)(my_view)
        response = view(request)
        self.assertEqual(response['Content-Security-Policy'], middleware._build_policy(policy))

    def test_csp_exclude(self):
        # Create a request object with a method that is exempt from CSP
        request = self.factory.get('/my-view/')

        # Create a response object with the CSP header
        response = HttpResponse()
        middleware = CSPMiddleware(get_response=None)
        response['Content-Security-Policy'] = middleware._build_policy({})

        # Set the exempt_urls attribute to include the URL path of the request object
        middleware.exempt_urls = ['/my-view/']

        # Verify that the process_request method returns None since the URL path of the request is exempt from CSP
        self.assertIsNone(middleware.process_request(request))

from functools import wraps
from django_csp.utils import get_csp_from_request

def csp_exempt(methods=None):
    """
    Decorator to exempt a view from CSP enforcement for specific request methods.
    """
    if not methods:
        methods = []

    def decorator(view_func):
        @wraps(view_func)
        def _wrapped_view(request, *args, **kwargs):
            if request.method in methods:
                # Remove CSP from the request object
                request.csp = None
            return view_func(request, *args, **kwargs)

        return _wrapped_view

    return decorator

from functools import wraps
from django_csp.utils import get_csp_from_request

def csp_update(methods=None, **kwargs):
    """
    Decorator to update CSP directives for specific request methods.
    """
    if not methods:
        methods = []

    def decorator(view_func):
        @wraps(view_func)
        def _wrapped_view(request, *args, **kwargs):
            if request.method in methods:
                # Update the CSP for the current request method
                csp = get_csp_from_request(request)
                for directive, values in kwargs.items():
                    csp[directive] = values
            return view_func(request, *args, **kwargs)

        return _wrapped_view

    return decorator

from django.http import HttpResponse
from django.views import View

from csp_decorators import csp_exempt, csp_update


@csp_exempt(['GET', 'HEAD'])
@csp_update(['POST'], script_src=['https://example.com/])
class MyView(View):
    def get(self, request, *args, kwargs):
        content = "This is a GET request"
        return HttpResponse(content)

    def post(self, request, *args, **kwargs):
        content = "This is a POST request"
        return HttpResponse(content)

from django.test import RequestFactory, TestCase
from django.http import HttpResponse
from django.urls import reverse
from django_csp.middleware import CSPMiddleware
from myapp.views import MyView

class CSPDecoratorsTest(TestCase):
    def setUp(self):
        self.factory = RequestFactory()

    def test_csp_exempt(self):
        # Create a GET request object and pass it through the csp_exempt decorator, which should remove the CSP from the request object
        get_request = self.factory.get(reverse('my-view'))
        get_response = MyView.as_view()(get_request)
        self.assertEqual(get_response.status_code, 200)
        self.assertFalse(hasattr(get_request, 'csp'))

        # Create a POST request object and pass it through the csp_exempt decorator, which should remove the CSP from the request object
        post_request = self.factory.post(reverse('my-view'))
        post_response = MyView.as_view()(post_request)
        self.assertEqual(post_response.status_code, 200)
        self.assertFalse(hasattr(post_request, 'csp'))

    def test_csp_update(self):
        # Create a POST request object and pass it through the csp_update decorator, which should add a nonce to the CSP of the request object
        post_request = self.factory.post(reverse('my-view'))
        post_response = MyView.as_view()(post_request)
        self.assertEqual(post_response.status_code, 200)
        self.assertTrue(hasattr(post_request, 'csp'))
        self.assertIn('nonce-', post_request.csp)

from functools import wraps

def valid_method(allowed_methods, request):
    """
        Params:
            allowed_methods:    list or string of method to make the update
            request:            the request for the csp

        Returns:
            True if the method in decorator parameter matches

        Check if any method restriction is available. If not, update is valid.
        If a list of methods is given, iterate through methods to check if request method is in
        restriction
        If only a method is given, check if the request method is the matching one
    """
    if allowed_methods is None:
        return True
    if isinstance(request, tuple):
        request = request[0]
    result = False
    if isinstance(allowed_methods, list):
        for method in allowed_methods:
            if request.method == method:
                result = True
                break
    elif isinstance(allowed_methods, str):
        if request.method == allowed_methods:
            result = True
    return result

def csp_exempt(methods=None, **kwargs):
    """
        csp_exempt
    """
    def decorator(function):
        @wraps(function)
        def _wrapped(*a, **kw):
            response = function(*a, **kw)
            if valid_method(methods, a):
                response._csp_exempt = True
            return response
        return _wrapped
    return decorator


def csp_update(methods=None, **kwargs):
    """
        csp_update
    """
    update = dict((k.lower().replace('_', '-'), v) for k, v in kwargs.items())

    def decorator(function):
        @wraps(function)
        def _wrapped(*a, **kw):
            response = function(*a, **kw)
            if valid_method(methods, a):
                response._csp_update = update
            return response
        return _wrapped
    return decorator


def csp_replace(methods=None, **kwargs):
    """
        csp_replace
    """
    replace = dict((k.lower().replace('_', '-'), v) for k, v in kwargs.items())

    def decorator(function):
        @wraps(function)
        def _wrapped(*a, **kw):
            response = function(*a, **kw)
            if valid_method(methods, a):
                response._csp_replace = replace
            return response
        return _wrapped
    return decorator


def csp(methods=None, **kwargs):
    """
        csp
    """
    config = dict(
        (k.lower().replace('_', '-'), [v] if isinstance(v, str) else v)
        for k, v
        in kwargs.items()
    )

    def decorator(function):
        @wraps(function)
        def _wrapped(*a, **kw):
            response = function(*a, **kw)
            if valid_method(methods, a):
                response._csp_config = config
            return response
        return _wrapped
    return decorator