diff --git a/dojo/engagement/views.py b/dojo/engagement/views.py index ea73bd80c63..b813a9c2758 100644 --- a/dojo/engagement/views.py +++ b/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") diff --git a/dojo/forms.py b/dojo/forms.py index 6fe83668d1b..cb1e670054e 100644 --- a/dojo/forms.py +++ b/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")) diff --git a/dojo/utils.py b/dojo/utils.py index 470d8607725..8bbd5312107 100644 --- a/dojo/utils.py +++ b/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