FileUploads: Clean up media when related objects are deleted (#13028)

* FileUploads: Clean up `media` when related objects are deleted

* Update dojo/db_migrations/0241_file_upload_cleanup.py

* Handle case where the UI will attempt to delete files as well

* Update migrations

Author: Maffooch

dojo/apps.py

@@ -76,6 +76,7 @@ def ready(self):
         import dojo.cred.signals
         import dojo.endpoint.signals
         import dojo.engagement.signals
+        import dojo.file_uploads.signals
         import dojo.finding_group.signals
         import dojo.notes.signals
         import dojo.product.signals

dojo/db_migrations/0242_file_upload_cleanup.py

@@ -0,0 +1,39 @@
+from django.db import migrations
+from django.db.models import Q
+
+
+def delete_unreferenced_file_uploads(apps, schema_editor):
+    FileUpload = apps.get_model("dojo", "FileUpload")
+    # Filter files that have no relations to Finding, Test, or Engagement
+    unused_files = FileUpload.objects.filter(
+        Q(finding__isnull=True) &
+        Q(test__isnull=True) &
+        Q(engagement__isnull=True)
+    ).distinct()
+    # Delete the files from disk first, then delete the FileUpload object
+    for file_upload in unused_files:
+        if file_upload.file:
+            storage = file_upload.file.storage
+            path = file_upload.file.path
+            if path and storage.exists(path):
+                storage.delete(path)
+        file_upload.delete()
+
+
+def cannot_turn_back_time(apps, schema_editor):
+    """
+    We cannot possibly return to the original state without knowing
+    the original value at the time the migration is revoked. Instead
+    we will do nothing.
+    """
+    pass
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("dojo", "0241_remove_system_settings_time_zone"), 
+    ]
+
+    operations = [
+        migrations.RunPython(delete_unreferenced_file_uploads, cannot_turn_back_time),
+    ]

dojo/engagement/signals.py

@@ -8,6 +8,7 @@
 from django.urls import reverse
 from django.utils.translation import gettext as _
 
+from dojo.file_uploads.helper import delete_related_files
 from dojo.models import Engagement, Product
 from dojo.notes.helper import delete_related_notes
 from dojo.notifications.helper import create_notification
@@ -65,3 +66,4 @@ def engagement_post_delete(sender, instance, using, origin, **kwargs):
 def engagement_pre_delete(sender, instance, **kwargs):
     with contextlib.suppress(sender.DoesNotExist):
         delete_related_notes(instance)
+        delete_related_files(instance)

dojo/file_uploads/__init__.py

@@ -0,0 +1,11 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def delete_related_files(obj):
+    if not hasattr(obj, "files"):
+        logger.warning(f"Attempted to delete files from object type {type(obj)} without 'files' attribute.")
+        return
+    logger.debug(f"Deleting {obj.files.count()} files for {type(obj).__name__} {obj.id}")
+    obj.files.all().delete()

dojo/file_uploads/helper.py

@@ -0,0 +1,35 @@
+from django.db.models.signals import post_delete, pre_save
+from django.dispatch import receiver
+
+from dojo.models import FileUpload
+
+
+@receiver(post_delete, sender=FileUpload)
+def delete_file_on_object_delete(sender, instance, **kwargs):
+    """
+    Deletes file from filesystem when corresponding `FileUpload` is deleted.
+    Mostly used as a backup in case the FileUpload.delete() function fails
+    for whatever reason
+    """
+    if instance.file:
+        storage = instance.file.storage
+        path = instance.file.path
+        if path and storage.exists(path):
+            storage.delete(path)
+
+
+@receiver(pre_save, sender=FileUpload)
+def delete_old_file_on_change(sender, instance, **kwargs):
+    """Deletes old file when a new file is uploaded to the same record."""
+    if not instance.pk:
+        return  # Skip new objects
+    try:
+        old_file = FileUpload.objects.get(pk=instance.pk).file
+    except FileUpload.DoesNotExist:
+        return
+    new_file = instance.file
+    if old_file and old_file != new_file:
+        storage = old_file.storage
+        path = old_file.path
+        if path and storage.exists(path):
+            storage.delete(path)

dojo/file_uploads/signals.py

@@ -14,6 +14,7 @@
 from dojo.celery import app
 from dojo.decorators import dojo_async_task, dojo_model_from_id, dojo_model_to_id
 from dojo.endpoint.utils import save_endpoints_to_add
+from dojo.file_uploads.helper import delete_related_files
 from dojo.models import (
     Endpoint,
     Endpoint_Status,
@@ -413,6 +414,7 @@ def finding_pre_delete(sender, instance, **kwargs):
     # https://code.djangoproject.com/ticket/154
     instance.found_by.clear()
     delete_related_notes(instance)
+    delete_related_files(instance)
 
 
 def finding_delete(instance, **kwargs):

dojo/finding/helper.py

@@ -790,6 +790,14 @@ class FileUpload(models.Model):
     title = models.CharField(max_length=100, unique=True)
     file = models.FileField(upload_to=UniqueUploadNameProvider("uploaded_files"))
 
+    def delete(self, *args, **kwargs):
+        """Delete the model and remove the file from storage."""
+        storage = self.file.storage
+        path = self.file.path
+        super().delete(*args, **kwargs)
+        if path and storage.exists(path):
+            storage.delete(path)
+
     def copy(self):
         copy = _copy_model_util(self)
         # Add unique modifier to file name

dojo/models.py

@@ -8,6 +8,7 @@
 from django.urls import reverse
 from django.utils.translation import gettext as _
 
+from dojo.file_uploads.helper import delete_related_files
 from dojo.models import Engagement, Finding, Product, Test
 from dojo.notes.helper import delete_related_notes
 from dojo.notifications.helper import create_notification
@@ -57,3 +58,4 @@ def update_found_by_for_findings(sender, instance, **kwargs):
 def test_pre_delete(sender, instance, **kwargs):
     with contextlib.suppress(sender.DoesNotExist):
         delete_related_notes(instance)
+        delete_related_files(instance)

dojo/test/signals.py

@@ -1,4 +1,5 @@
 import logging
+from contextlib import suppress
 from pathlib import Path
 
 from auditlog.models import LogEntry
@@ -157,7 +158,8 @@ def manage_files(request, oid, obj_type):
 
             for o in files_formset.deleted_objects:
                 logger.debug("removing file: %s", o.file.name)
-                (Path(settings.MEDIA_ROOT) / o.file.name).unlink()
+                with suppress(FileNotFoundError):
+                    (Path(settings.MEDIA_ROOT) / o.file.name).unlink()
 
             for o in files_formset.new_objects:
                 logger.debug("adding file: %s", o.file.name)
@@ -168,7 +170,8 @@ def manage_files(request, oid, obj_type):
                                                      finding__isnull=True)
             for o in orphan_files:
                 logger.debug("purging orphan file: %s", o.file.name)
-                (Path(settings.MEDIA_ROOT) / o.file.name).unlink()
+                with suppress(FileNotFoundError):
+                    (Path(settings.MEDIA_ROOT) / o.file.name).unlink()
                 o.delete()
 
             messages.add_message(