diff --git a/atst/domain/csp/files.py b/atst/domain/csp/files.py index dfb48a3c..74384500 100644 --- a/atst/domain/csp/files.py +++ b/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 diff --git a/atst/forms/task_order.py b/atst/forms/task_order.py index 0c87c42c..f4e21baf 100644 --- a/atst/forms/task_order.py +++ b/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"}, ) diff --git a/atst/forms/validators.py b/atst/forms/validators.py index ff4b2366..99265769 100644 --- a/atst/forms/validators.py +++ b/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 diff --git a/tests/forms/test_validators.py b/tests/forms/test_validators.py index 2dc80659..c57fe5e0 100644 --- a/tests/forms/test_validators.py +++ b/tests/forms/test_validators.py @@ -1,13 +1,7 @@ from wtforms.validators import ValidationError, StopValidation import pytest -from atst.forms.validators import ( - Name, - IsNumber, - PhoneNumber, - ListItemsUnique, - RequiredIf, -) +from atst.forms.validators import * 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)