From 7507fa50b36e6fe87c6a4e3f77aa7926fbcb83cb Mon Sep 17 00:00:00 2001 From: Lacey Williams Henschel Date: Thu, 2 Nov 2023 11:45:52 -0700 Subject: [PATCH] Limit file types of user and news images - Add file type validator - Apply validator to user image - Save validator class to variable - Add image validator to news image model field - Add file size validator - Enable file size validator on news and user model image fields - Fix test --- core/tests/test_validators.py | 53 +++++++++++++++++++++++ core/validators.py | 44 +++++++++++++++++++ news/migrations/0008_alter_entry_image.py | 28 ++++++++++++ news/models.py | 9 +++- news/tests/test_models.py | 50 +++++++++++++++++++++ news/tests/test_views.py | 2 +- users/migrations/0011_alter_user_image.py | 28 ++++++++++++ users/models.py | 9 +++- users/tests/test_models.py | 47 +++++++++++++++++--- 9 files changed, 261 insertions(+), 9 deletions(-) create mode 100644 core/tests/test_validators.py create mode 100644 core/validators.py create mode 100644 news/migrations/0008_alter_entry_image.py create mode 100644 users/migrations/0011_alter_user_image.py diff --git a/core/tests/test_validators.py b/core/tests/test_validators.py new file mode 100644 index 00000000..225412ec --- /dev/null +++ b/core/tests/test_validators.py @@ -0,0 +1,53 @@ +import pytest + +from django.core.exceptions import ValidationError +from django.core.files.uploadedfile import SimpleUploadedFile +from ..validators import FileTypeValidator, MaxFileSizeValidator + + +def test_file_type_validator(): + """ + Eensure validator only allows specific file extensions. + """ + validator = FileTypeValidator(extensions=[".jpg", ".png"]) + + # Valid file types + valid_file = SimpleUploadedFile( + "test.jpg", b"file_content", content_type="image/jpeg" + ) + validator(valid_file) + + valid_file = SimpleUploadedFile( + "test.png", b"file_content", content_type="image/jpeg" + ) + validator(valid_file) + + # Invalid file type + invalid_file = SimpleUploadedFile( + "test.txt", b"file_content", content_type="text/plain" + ) + with pytest.raises(ValidationError): + validator(invalid_file) + + +def test_max_file_size_validator(): + """ + Ensure the validator enforces the file size limit. + """ + # 1MB max file size + validator = MaxFileSizeValidator(max_size=1 * 1024 * 1024) + + # Valid file size + valid_file = SimpleUploadedFile( + "small.jpg", b"a" * (1 * 1024 * 1024 - 1), content_type="image/jpeg" + ) + validator(valid_file) # Should not raise + + # Invalid file size + invalid_file = SimpleUploadedFile( + "large.jpg", b"a" * (1 * 1024 * 1024 + 1), content_type="image/jpeg" + ) + with pytest.raises(ValidationError) as exc_info: + validator(invalid_file) + + assert "File too large" in str(exc_info.value) diff --git a/core/validators.py b/core/validators.py new file mode 100644 index 00000000..56654df2 --- /dev/null +++ b/core/validators.py @@ -0,0 +1,44 @@ +import os + +from django.core.exceptions import ValidationError +from django.utils.deconstruct import deconstructible + + +@deconstructible +class FileTypeValidator: + """ + Validates that a file is of the permitted types. + """ + + def __init__(self, extensions): + self.extensions = extensions + + def __call__(self, value): + ext = os.path.splitext(value.name)[1].lower() + if ext not in self.extensions: + raise ValidationError( + f"Unsupported file extension. Allowed types are: {', '.join(self.extensions)}" # noqa + ) + + +image_validator = FileTypeValidator(extensions=[".jpg", ".jpeg", ".png"]) + + +@deconstructible +class MaxFileSizeValidator: + """ + Validates that a file is not larger than a certain size. + """ + + def __init__(self, max_size): + self.max_size = max_size + + def __call__(self, value): + if value.size > self.max_size: + raise ValidationError( + f"File too large. Size should not exceed {self.max_size / 1024 / 1024} MB." # noqa + ) + + +# 1 MB max file size +max_file_size_validator = MaxFileSizeValidator(max_size=1 * 1024 * 1024) diff --git a/news/migrations/0008_alter_entry_image.py b/news/migrations/0008_alter_entry_image.py new file mode 100644 index 00000000..b9a779fe --- /dev/null +++ b/news/migrations/0008_alter_entry_image.py @@ -0,0 +1,28 @@ +# Generated by Django 4.2.2 on 2023-11-02 19:32 + +import core.validators +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("news", "0007_alter_entry_external_url_alter_entry_slug"), + ] + + operations = [ + migrations.AlterField( + model_name="entry", + name="image", + field=models.ImageField( + blank=True, + null=True, + upload_to="news/%Y/%m/", + validators=[ + core.validators.FileTypeValidator( + extensions=[".jpg", ".jpeg", ".png"] + ), + core.validators.MaxFileSizeValidator(max_size=1048576), + ], + ), + ), + ] diff --git a/news/models.py b/news/models.py index 002982b8..75c5789d 100644 --- a/news/models.py +++ b/news/models.py @@ -7,6 +7,8 @@ from django.utils.text import slugify from django.utils.timezone import now from django.utils.translation import gettext_lazy as _ +from core.validators import image_validator, max_file_size_validator + from . import acl User = get_user_model() @@ -67,7 +69,12 @@ class Entry(models.Model): related_name="moderated_entries_set", ) external_url = models.URLField(_("URL"), blank=True, default="", max_length=500) - image = models.ImageField(upload_to="news/%Y/%m/", null=True, blank=True) + image = models.ImageField( + upload_to="news/%Y/%m/", + null=True, + blank=True, + validators=[image_validator, max_file_size_validator], + ) created_at = models.DateTimeField(default=now) approved_at = models.DateTimeField(null=True, blank=True) modified_at = models.DateTimeField(auto_now=True) diff --git a/news/tests/test_models.py b/news/tests/test_models.py index 5ff7c06c..d3037aa3 100644 --- a/news/tests/test_models.py +++ b/news/tests/test_models.py @@ -2,6 +2,8 @@ import datetime import pytest from django.contrib.auth.models import Permission +from django.core.files.uploadedfile import SimpleUploadedFile +from django.core.exceptions import ValidationError from django.utils.timezone import now from model_bakery import baker @@ -62,6 +64,54 @@ def test_entry_absolute_url(tp): assert entry.get_absolute_url() == tp.reverse("news-detail", "the-slug") +def test_entry_model_image_validator(tp): + """ + Test that the `image` field only accepts certain file types. + """ + author = baker.make("users.User") + entry = Entry.objects.create(title="😀 Foo Bar Baz!@! +", author=author) + # Valid image file + valid_image = SimpleUploadedFile( + "test.jpg", b"file_content", content_type="image/jpeg" + ) + entry.image = valid_image + # This should not raise any errors + entry.full_clean() + + # Invalid image file + invalid_image = SimpleUploadedFile( + "test.pdf", b"file_content", content_type="application/pdf" + ) + entry.image = invalid_image + # This should raise a ValidationError + with pytest.raises(ValidationError): + entry.full_clean() + + +def test_entry_model_image_file_size(tp): + """ + Test that the `image` field rejects files larger than a specific size. + """ + author = baker.make("users.User") + entry = Entry.objects.create(title="😀 Foo Bar Baz!@! +", author=author) + + valid_image = SimpleUploadedFile( + "test.jpg", b"a" * (1 * 1024 * 1024 - 1), content_type="image/jpeg" + ) + entry.image = valid_image + # This should not raise any errors + entry.full_clean() + + # This should fail (just over 1MB) + invalid_image = SimpleUploadedFile( + "too_large.jpg", b"a" * (1 * 1024 * 1024 + 1), content_type="image/jpeg" + ) + entry.image = invalid_image + # This should raise a ValidationError for file size + with pytest.raises(ValidationError): + entry.full_clean() + + def test_approve_entry(make_entry): future = now() + datetime.timedelta(hours=1) entry = make_entry(approved=False, publish_at=future) diff --git a/news/tests/test_views.py b/news/tests/test_views.py index 8484b395..f3557fa6 100644 --- a/news/tests/test_views.py +++ b/news/tests/test_views.py @@ -418,7 +418,7 @@ def test_news_create_post( b"GIF89a\x01\x00\x01\x00\x00\x00\x00!\xf9\x04\x01\x00\x00\x00" b"\x00,\x00\x00\x00\x00\x01\x00\x01\x00\x00\x02\x01\x00\x00" ) - img.name = f"random-value-{uuid.uuid4()}.gif" + img.name = f"random-value-{uuid.uuid4()}.png" data["image"] = img data["publish_at"] = right_now = now() diff --git a/users/migrations/0011_alter_user_image.py b/users/migrations/0011_alter_user_image.py new file mode 100644 index 00000000..328d92e4 --- /dev/null +++ b/users/migrations/0011_alter_user_image.py @@ -0,0 +1,28 @@ +# Generated by Django 4.2.2 on 2023-11-02 19:32 + +import core.validators +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("users", "0010_preferences"), + ] + + operations = [ + migrations.AlterField( + model_name="user", + name="image", + field=models.FileField( + blank=True, + null=True, + upload_to="profile-images", + validators=[ + core.validators.FileTypeValidator( + extensions=[".jpg", ".jpeg", ".png"] + ), + core.validators.MaxFileSizeValidator(max_size=1048576), + ], + ), + ), + ] diff --git a/users/models.py b/users/models.py index f0e1dbf6..24b3346c 100644 --- a/users/models.py +++ b/users/models.py @@ -16,6 +16,8 @@ from django.dispatch import receiver from django.utils import timezone from django.utils.translation import gettext_lazy as _ +from core.validators import image_validator, max_file_size_validator + logger = logging.getLogger(__name__) @@ -203,7 +205,12 @@ class User(BaseUser): badges = models.ManyToManyField(Badge) github_username = models.CharField(_("github username"), max_length=100, blank=True) - image = models.FileField(upload_to="profile-images", null=True, blank=True) + image = models.FileField( + upload_to="profile-images", + null=True, + blank=True, + validators=[image_validator, max_file_size_validator], + ) claimed = models.BooleanField( _("claimed"), default=True, diff --git a/users/tests/test_models.py b/users/tests/test_models.py index 8044923e..dbd965d7 100644 --- a/users/tests/test_models.py +++ b/users/tests/test_models.py @@ -1,6 +1,8 @@ import pytest from django.contrib.auth import get_user_model +from django.core.files.uploadedfile import SimpleUploadedFile +from django.core.exceptions import ValidationError from pytest_django.asserts import assertQuerySetEqual from ..models import Preferences @@ -50,14 +52,47 @@ def test_get_display_name(user): assert user.get_display_name == "Last" -@pytest.mark.skip("Add this test when I have the patience for mocks") -def test_user_save_image_from_github(user): +def test_user_model_image_validator(user): """ - Test `User.save_image_from_github(avatar_url)` - See test_signals -- you will need to do something similar here, but - dealing with a File object might make it trickier. + Test that the `image` field on the User model only accepts certain file types. """ - pass + # Valid image file + valid_image = SimpleUploadedFile( + "test.jpg", b"file_content", content_type="image/jpeg" + ) + user.image = valid_image + # This should not raise any errors + user.full_clean() + + # Invalid image file + invalid_image = SimpleUploadedFile( + "test.pdf", b"file_content", content_type="application/pdf" + ) + user.image = invalid_image + # This should raise a ValidationError + with pytest.raises(ValidationError): + user.full_clean() + + +def test_user_model_image_file_size(user): + """ + Test that the `image` field rejects files larger than a specific size. + """ + valid_image = SimpleUploadedFile( + "test.jpg", b"a" * (1 * 1024 * 1024 - 1), content_type="image/jpeg" + ) + user.image = valid_image + # This should not raise any errors + user.full_clean() + + # This should fail (just over 1MB) + invalid_image = SimpleUploadedFile( + "too_large.jpg", b"a" * (1 * 1024 * 1024 + 1), content_type="image/jpeg" + ) + user.image = invalid_image + # This should raise a ValidationError for file size + with pytest.raises(ValidationError): + user.full_clean() def test_claim(user):