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,53 @@ def test_banners_no_xss(self):
         self.assertNotIn("Welcome to Nautobot!", response.content.decode(response.charset))
 
 
+class MediaViewTestCase(TestCase):
+    def test_media_unauthenticated(self):
+        """
+        Test that unauthenticated users are redirected to the login page
+        when attempting to access non-branded media files.
+        """
+        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_branding_media_unauthenticated(self):
+        """
+        Test that unauthenticated users can access branded media files
+        listed in `settings.BRANDING_FILEPATHS`.
+        """
+        with tempfile.TemporaryDirectory() as temp_dir:
+            with override_settings(MEDIA_ROOT=temp_dir, BRANDING_FILEPATHS={"logo": os.path.join(temp_dir, "foo.txt")}):
+                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"})
+                self.client.logout()
+                response = self.client.get(url)
+                self.assertHttpStatus(response, 200)
+
+    def test_media_authenticated(self):
+        """
+        Test that authenticated users can access regular media files
+        stored in the `MEDIA_ROOT`.
+        """
+        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 (
@@ -132,6 +133,34 @@ def get(self, request, *args, **kwargs):
         return self.render_to_response(context)
 
 
+class MediaView(AccessMixin, View):
+    """
+    Serves static media files while enforcing login restrictions.
+
+    This view wraps Django's `serve()` function to ensure that access to
+    media files is restricted to authenticated users. Additionally, it
+    allows access to branded media files defined in
+    `settings.BRANDING_FILEPATHS`, even for unauthenticated users.
+    """
+
+    def get(self, request, path):
+        can_view = request.user.is_authenticated
+        with contextlib.suppress(FileNotFoundError):
+            for branding_filepath in settings.BRANDING_FILEPATHS.values():
+                if not branding_filepath:
+                    continue
+                branding_filepath = os.path.join(settings.MEDIA_ROOT, branding_filepath)
+                requested_path = os.path.join(settings.MEDIA_ROOT, path)
+                same_path = os.path.samefile(branding_filepath, requested_path)
+                if same_path:
+                    can_view = True
+                    break
+
+        if not can_view:
+            return self.handle_no_permission()
+        return serve(request, path, document_root=settings.MEDIA_ROOT)
+
+
 class WorkerStatusView(UserPassesTestMixin, TemplateView):
     template_name = "utilities/worker_status.html"