Add validation for file extension in AbstractDocument using WAGTAIL_DOCS_EXTENSIONS setting (#6423)

docs/reference/settings.rst

@@ -199,6 +199,16 @@ Specifies the MIME content type that will be returned for the given file extensi
 
 A list of MIME content types that will be shown inline in the browser (by serving the HTTP header ``Content-Disposition: inline``) rather than served as a download, when using the ``serve_view`` method. Defaults to ``application/pdf``.
 
+.. _wagtaildocs_extensions:
+
+.. code-block:: python
+
+  WAGTAILDOCS_EXTENSIONS = ['pdf', 'docx']
+
+A list of allowed document extensions that will be validated during document uploading. 
+If this isn't supplied all document extensions are allowed. 
+Warning: this doesn't always ensure that the uploaded file is valid as files can 
+be renamed to have an extension no matter what data they contain.
 
 Password Management
 ===================

wagtail/documents/models.py

@@ -5,6 +5,7 @@ from contextlib import contextmanager
 from mimetypes import guess_type
 
 from django.conf import settings
+from django.core.validators import FileExtensionValidator
 from django.db import models
 from django.dispatch import Signal
 from django.urls import reverse
@@ -53,6 +54,21 @@ class AbstractDocument(CollectionMember, index.Indexed, models.Model):
         index.FilterField('uploaded_by_user'),
     ]
 
+    def clean(self):
+        """
+        Checks for WAGTAILDOCS_EXTENSIONS and validates the uploaded file
+        based on allowed extensions that were specified.
+        Warning : This doesn't always ensure that the uploaded file is valid
+        as files can be renamed to have an extension no matter what
+        data they contain.
+
+        More info : https://docs.djangoproject.com/en/3.1/ref/validators/#fileextensionvalidator
+        """
+        allowed_extensions = getattr(settings, "WAGTAILDOCS_EXTENSIONS", None)
+        if allowed_extensions:
+            validate = FileExtensionValidator(allowed_extensions)
+            validate(self.file)
+
     def is_stored_locally(self):
         """
         Returns True if the image is hosted on the local filesystem

wagtail/documents/tests/test_models.py

@@ -1,6 +1,6 @@
 from django.conf import settings
 from django.contrib.auth.models import Group, Permission
-from django.core.exceptions import ImproperlyConfigured
+from django.core.exceptions import ImproperlyConfigured, ValidationError
 from django.core.files.base import ContentFile
 from django.db import transaction
 from django.test import TestCase, TransactionTestCase
@@ -161,6 +161,47 @@ class TestFilesDeletedForDefaultModels(TransactionTestCase):
         self.assertFalse(document.file.storage.exists(filename))
 
 
+@override_settings(WAGTAILDOCS_EXTENSIONS=["pdf"])
+class TestDocumentValidateExtensions(TestCase):
+    def setUp(self):
+        self.document_invalid = models.Document.objects.create(
+            title="Test document", file="test.doc"
+        )
+        self.document_valid = models.Document.objects.create(
+            title="Test document", file="test.pdf"
+        )
+
+    def test_create_doc_invalid_extension(self):
+        """
+        Checks if the uploaded document has the expected extensions
+        mentioned in settings.WAGTAILDOCS_EXTENSIONS
+
+        This is caught in form.error and should be raised during model
+        creation when called full_clean. This specific testcase invalid
+        file extension is passed
+        """
+        with self.assertRaises(ValidationError):
+            self.document_invalid.full_clean()
+
+    def test_create_doc_valid_extension(self):
+        """
+        Checks if the uploaded document has the expected extensions
+        mentioned in settings.WAGTAILDOCS_EXTENSIONS
+
+        This is caught in form.error and should be raised during
+        model creation when called full_clean. In this specific
+        testcase invalid file extension is passed.
+        """
+        try:
+            self.document_valid.full_clean()
+        except ValidationError:
+            self.fail("Validation error is raised even when valid file name is passed")
+
+    def tearDown(self):
+        self.document_invalid.file.delete()
+        self.document_valid.file.delete()
+
+
 @override_settings(WAGTAILDOCS_DOCUMENT_MODEL='tests.CustomDocument')
 class TestFilesDeletedForCustomModels(TestFilesDeletedForDefaultModels):
     def setUp(self):