File length validation for task order upload

atst/domain/csp/files.py

@@ -14,7 +14,7 @@ class CSPFileError(Exception):
 
 
 class FileProviderInterface:
-    _PERMITTED_MIMETYPES = ["application/pdf", "image/png"]
+    _PERMITTED_MIMETYPES = ["application/pdf"]
 
     def _enforce_mimetype(self, fyle):
         # TODO: for hardening, we should probably use a better library for

atst/forms/task_order.py

@@ -4,6 +4,7 @@ from wtforms.validators import Required, Optional
 from flask_wtf.file import FileAllowed
 
 from .forms import BaseForm
+from atst.forms.validators import FileLength
 from atst.utils.localization import translate
 
 
@@ -16,7 +17,8 @@ class TaskOrderForm(BaseForm):
     pdf = FileField(
         None,
         validators=[
-            FileAllowed(["pdf"], translate("forms.task_order.file_format_not_allowed"))
+            FileAllowed(["pdf"], translate("forms.task_order.file_format_not_allowed")),
+            FileLength(),
         ],
         render_kw={"accept": ".pdf,application/pdf"},
     )

atst/forms/validators.py

@@ -99,3 +99,17 @@ def RequiredIf(criteria_function, message=translate("forms.validators.is_require
             raise StopValidation()
 
     return _required_if
+
+
+def FileLength(max_length=50000000, message=None):
+    def _file_length(_form, field):
+        if field.data is None:
+            return True
+
+        content = field.data.read()
+        if len(content) > max_length:
+            raise ValidationError(message)
+        else:
+            field.data.seek(0)
+
+    return _file_length

tests/forms/test_validators.py

@@ -1,13 +1,7 @@
 from wtforms.validators import ValidationError, StopValidation
 import pytest
 
-from atst.forms.validators import (
+from atst.forms.validators import *
-    Name,
-    IsNumber,
-    PhoneNumber,
-    ListItemsUnique,
-    RequiredIf,
-)
 
 
 class TestIsNumber:
@@ -97,3 +91,12 @@ class TestRequiredIf:
 
         with pytest.raises(StopValidation):
             validator(dummy_form, dummy_field)
+
+
+class TestFileLength:
+    def test_FileLength(self, dummy_form, dummy_field, pdf_upload):
+        validator = FileLength(max_length=1)
+        dummy_field.data = pdf_upload
+
+        with pytest.raises(ValidationError):
+            validator(dummy_form, dummy_field)