nautobot · timizuoebideri1 · Dec 17, 2024 · Dec 19, 2024 · Dec 19, 2024 · Dec 30, 2024
diff --git a/nautobot/core/tests/test_views.py b/nautobot/core/tests/test_views.py
@@ -1,4 +1,6 @@
+import os
 import re
+import tempfile
 from unittest import mock, skipIf
 import urllib.parse
 
@@ -181,6 +183,29 @@ def test_banners_no_xss(self):
         self.assertNotIn("Welcome to Nautobot!", response.content.decode(response.charset))
 
 
+class MediaViewTestCase(TestCase):
+    def test_media_unauthenticated(self):
+        url = reverse("media", kwargs={"path": "foo.txt"})
+        self.client.logout()
+        response = self.client.get(url)
+
+        # Unauthenticated request should redirect to login page
+        self.assertRedirects(
+            response, expected_url=reverse("login") + "?next=/media/foo.txt", status_code=302, target_status_code=200
+        )
+
+    def test_media_authenticated(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            with override_settings(MEDIA_ROOT=temp_dir):
+                file_path = os.path.join(temp_dir, "foo.txt")
+                with open(file_path, "w") as f:
+                    f.write("Hello, world!")
+
+                url = reverse("media", kwargs={"path": "foo.txt"})
+                response = self.client.get(url)
+                self.assertHttpStatus(response, 200)
+
+
 @override_settings(BRANDING_TITLE="Nautobot")
 class SearchFieldsTestCase(TestCase):
     def test_search_bar_redirect_to_login(self):

diff --git a/nautobot/core/urls.py b/nautobot/core/urls.py
@@ -1,12 +1,12 @@
 from django.conf import settings
 from django.urls import include, path
 from django.views.generic import TemplateView
-from django.views.static import serve
 
 from nautobot.core.views import (
     CustomGraphQLView,
     get_file_with_authorization,
     HomeView,
+    MediaView,
     NautobotMetricsView,
     NautobotMetricsViewAuth,
     SearchView,
@@ -46,7 +46,7 @@
     # GraphQL
     path("graphql/", CustomGraphQLView.as_view(graphiql=True), name="graphql"),
     # Serving static media in Django (TODO: should be DEBUG mode only - "This view is NOT hardened for production use")
-    path("media/<path:path>", serve, {"document_root": settings.MEDIA_ROOT}),
+    path("media/<path:path>", MediaView.as_view(), name="media"),
     # Admin
     path("admin/", admin_site.urls),
     # Errors

diff --git a/nautobot/core/views/__init__.py b/nautobot/core/views/__init__.py
@@ -23,6 +23,7 @@
 from django.views.decorators.csrf import requires_csrf_token
 from django.views.defaults import ERROR_500_TEMPLATE_NAME, page_not_found
 from django.views.generic import TemplateView, View
+from django.views.static import serve
 from graphene_django.views import GraphQLView
 from packaging import version
 from prometheus_client import (
@@ -47,6 +48,7 @@
 from nautobot.core.releases import get_latest_release
 from nautobot.core.utils.lookup import get_route_for_model
 from nautobot.core.utils.permissions import get_permission_for_model
+from nautobot.core.views.generic import GenericView
 from nautobot.extras.forms import GraphQLQueryForm
 from nautobot.extras.models import FileProxy, GraphQLQuery, Status
 from nautobot.extras.registry import registry
@@ -132,6 +134,15 @@ def get(self, request, *args, **kwargs):
         return self.render_to_response(context)
 
 
+class MediaView(GenericView):
+    """
+    Wrap Django's serve() view to enforce LOGIN_REQUIRED for static media.
+    """
+
+    def get(self, request, path):
+        return serve(request, path, document_root=settings.MEDIA_ROOT)
+
+
 class WorkerStatusView(UserPassesTestMixin, TemplateView):
     template_name = "utilities/worker_status.html"