Threat Uploads: Server side file extension validation + force downloads

dojo/engagement/views.py

@@ -17,7 +17,7 @@
 from django.db import DEFAULT_DB_ALIAS
 from django.db.models import Count, Q
 from django.db.models.query import Prefetch, QuerySet
-from django.http import FileResponse, HttpRequest, HttpResponse, HttpResponseRedirect, QueryDict, StreamingHttpResponse
+from django.http import HttpRequest, HttpResponse, HttpResponseRedirect, QueryDict, StreamingHttpResponse
 from django.shortcuts import get_object_or_404, render
 from django.urls import Resolver404, reverse
 from django.utils import timezone
@@ -100,6 +100,7 @@
     add_success_message_to_response,
     async_delete,
     calculate_grade,
+    generate_file_response_from_file_path,
     get_cal_event,
     get_page_items,
     get_return_url,
@@ -1516,7 +1517,7 @@ def upload_threatmodel(request, eid):
 @user_is_authorized(Engagement, Permissions.Engagement_View, "eid")
 def view_threatmodel(request, eid):
     eng = get_object_or_404(Engagement, pk=eid)
-    return FileResponse(open(eng.tmodel_path, "rb"))
+    return generate_file_response_from_file_path(eng.tmodel_path)
 
 
 @user_is_authorized(Engagement, Permissions.Engagement_View, "eid")

dojo/forms.py

@@ -752,6 +752,23 @@ class UploadThreatForm(forms.Form):
         attrs={"accept": ".jpg,.png,.pdf"}),
         label="Select Threat Model")
 
+    def clean(self):
+        if (file := self.cleaned_data.get("file", None)) is not None:
+            ext = os.path.splitext(file.name)[1]  # [0] returns path+filename
+            valid_extensions = [".jpg", ".png", ".pdf"]
+            if ext.lower() not in valid_extensions:
+                if accepted_extensions := f"{', '.join(valid_extensions)}":
+                    msg = (
+                        "Unsupported extension. Supported extensions are as "
+                        f"follows: {accepted_extensions}"
+                    )
+                else:
+                    msg = (
+                        "File uploads are prohibited due to the list of acceptable "
+                        "file extensions being empty"
+                    )
+                raise ValidationError(msg)
+
 
 class MergeFindings(forms.ModelForm):
     FINDING_ACTION = (("", "Select an Action"), ("inactive", "Inactive"), ("delete", "Delete"))

dojo/utils.py

@@ -5,6 +5,7 @@
 import logging
 import mimetypes
 import os
+import pathlib
 import re
 from calendar import monthrange
 from datetime import date, datetime, timedelta
@@ -2616,14 +2617,32 @@ def generate_file_response(file_object: FileUpload) -> FileResponse:
         raise TypeError(msg)
     # Determine the path of the file on disk within the MEDIA_ROOT
     file_path = f"{settings.MEDIA_ROOT}/{file_object.file.url.lstrip(settings.MEDIA_URL)}"
-    _, file_extension = os.path.splitext(file_path)
+
+    return generate_file_response_from_file_path(
+        file_path, file_name=file_object.title, file_size=file_object.file.size,
+    )
+
+
+def generate_file_response_from_file_path(
+    file_path: str, file_name: str | None = None, file_size: int | None = None,
+) -> FileResponse:
+    """Serve an local file in a uniformed way."""
+    # Determine the file path
+    file_path_without_extension, file_extension = os.path.splitext(file_path)
+    # Determine the file name if not supplied
+    if file_name is None:
+        file_name = file_path_without_extension.rsplit("/")[-1]
+    # Determine the file size if not supplied
+    if file_size is None:
+        file_size = pathlib.Path(file_path).stat().st_size
     # Generate the FileResponse
+    full_file_name = f"{file_name}{file_extension}"
     response = FileResponse(
         open(file_path, "rb"),
-        filename=f"{file_object.title}{file_extension}",
+        filename=full_file_name,
         content_type=f"{mimetypes.guess_type(file_path)}",
     )
     # Add some important headers
-    response["Content-Disposition"] = f'attachment; filename="{file_object.title}{file_extension}"'
-    response["Content-Length"] = file_object.file.size
+    response["Content-Disposition"] = f'attachment; filename="{full_file_name}"'
+    response["Content-Length"] = file_size
     return response