Repository: mozmeao/snippets-service

Issue:
Enforce PNG and WebP formats for images



Discussion:

Resolving patch:
diff --git a/snippets/base/migrations/0001_initial.py b/snippets/base/migrations/0001_initial.py
--- a/snippets/base/migrations/0001_initial.py
+++ b/snippets/base/migrations/0001_initial.py
@@ -129,7 +129,7 @@ class Migration(migrations.Migration):
             fields=[
                 ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                 ('name', models.CharField(unique=True, max_length=255)),
-                ('code', models.TextField(validators=[snippets.base.models.validate_xml_template])),
+                ('code', models.TextField(validators=[snippets.base.validators.validate_xml_template])),
                 ('created', models.DateTimeField(auto_now_add=True)),
                 ('modified', models.DateTimeField(auto_now=True)),
             ],
diff --git a/snippets/base/models.py b/snippets/base/models.py
--- a/snippets/base/models.py
+++ b/snippets/base/models.py
@@ -26,7 +26,7 @@
 from snippets.base import util
 from snippets.base.fields import RegexField
 from snippets.base import managers
-from snippets.base.validators import validate_xml_template
+from snippets.base import validators
 
 
 JINJA_ENV = engines['backend']
@@ -82,7 +82,7 @@ class SnippetTemplate(models.Model):
         verbose_name='Priority template', default=False,
         help_text='Set to true to display first in dropdowns for faster selections')
     hidden = models.BooleanField(help_text='Hide from template selection dropdown', default=False)
-    code = models.TextField(validators=[validate_xml_template])
+    code = models.TextField(validators=[validators.validate_xml_template])
 
     created = models.DateTimeField(auto_now_add=True)
     modified = models.DateTimeField(auto_now=True)
@@ -560,7 +560,8 @@ class Icon(models.Model):
         upload_to=_generate_filename,
         height_field='height',
         width_field='width',
-        help_text=('PNGs only. Note that updating the image will '
+        validators=[validators.validate_image_format],
+        help_text=('PNG and WebP only. Note that updating the image will '
                    'update all snippets using this image.'),
     )
 
diff --git a/snippets/base/validators.py b/snippets/base/validators.py
--- a/snippets/base/validators.py
+++ b/snippets/base/validators.py
@@ -10,6 +10,7 @@
 from django.utils.deconstruct import deconstructible
 
 import bleach
+from PIL import Image
 
 ALLOWED_TAGS = ['a', 'i', 'b', 'u', 'strong', 'em', 'br']
 ALLOWED_ATTRIBUTES = {'a': ['href', 'data-metric']}
@@ -102,3 +103,10 @@ def validate_regex(regex_str):
         except re.error as exp:
             raise ValidationError(str(exp))
     return regex_str
+
+
+def validate_image_format(image):
+    img = Image.open(image)
+    if img.format not in ['PNG', 'WEBP']:
+        raise ValidationError('Upload only PNG or WebP images.')
+    return image


Test patch:
diff --git a/snippets/base/tests/test_validators.py b/snippets/base/tests/test_validators.py
--- a/snippets/base/tests/test_validators.py
+++ b/snippets/base/tests/test_validators.py
@@ -1,11 +1,14 @@
+import io
 import json
 from unittest.mock import patch
 
 from django.core.exceptions import ValidationError
 
+from PIL import Image
+
 from snippets.base.validators import (validate_as_router_fluent_variables,
                                       validate_xml_template, validate_xml_variables,
-                                      validate_regex)
+                                      validate_regex, validate_image_format)
 from snippets.base.tests import TestCase
 
 
@@ -75,3 +78,26 @@ def test_valid_regex(self):
     def test_invalid_regex(self):
         bogus_regex = r'/(?P\d+)/'
         self.assertRaises(ValidationError, validate_regex, bogus_regex)
+
+
+class ImageFormatValidatorTests(TestCase):
+    def test_valid_image_png(self):
+        img = Image.new('RGB', (30, 30), color='red')
+        fle = io.BytesIO()
+        img.save(fle, 'PNG')
+
+        self.assertEqual(validate_image_format(fle), fle)
+
+    def test_valid_image_webp(self):
+        img = Image.new('RGB', (30, 30), color='red')
+        fle = io.BytesIO()
+        img.save(fle, 'WEBP')
+
+        self.assertEqual(validate_image_format(fle), fle)
+
+    def test_invalid_image(self):
+        img = Image.new('RGB', (30, 30), color='red')
+        fle = io.BytesIO()
+        img.save(fle, 'JPEG')
+
+        self.assertRaises(ValidationError, validate_image_format, fle)