51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
"""Image uploads are decoded and re-encoded before they leave the application."""
|
|
from io import BytesIO
|
|
import unittest
|
|
|
|
from PIL import Image
|
|
|
|
from safe_images import MAX_SCREENSHOT_BYTES, ScreenshotError, sanitize_screenshot
|
|
|
|
|
|
class SafeScreenshotTests(unittest.TestCase):
|
|
def png(self, size=(32, 24), exif=None):
|
|
image = Image.new('RGB', size, (32, 64, 128))
|
|
out = BytesIO()
|
|
image.save(out, format='PNG', exif=exif or b'')
|
|
return out.getvalue()
|
|
|
|
def test_valid_raster_is_reencoded_without_metadata_or_appended_payload(self):
|
|
exif = Image.Exif()
|
|
exif[270] = 'private camera metadata fixture'
|
|
source = self.png(exif=exif)
|
|
payload = b'<script>not part of screenshot</script>'
|
|
cleaned, extension, media_type = sanitize_screenshot(source + payload)
|
|
self.assertEqual((extension, media_type), ('jpg','image/jpeg'))
|
|
self.assertTrue(cleaned.startswith(bytes((0xff,0xd8))))
|
|
self.assertNotIn(b'private camera metadata fixture',cleaned)
|
|
self.assertNotIn(payload,cleaned)
|
|
with Image.open(BytesIO(cleaned)) as image:
|
|
self.assertEqual(image.size,(32,24))
|
|
|
|
def test_alpha_image_remains_png(self):
|
|
image=Image.new('RGBA',(10,12),(0,0,0,0)); out=BytesIO(); image.save(out,format='PNG')
|
|
cleaned,extension,media_type=sanitize_screenshot(out.getvalue())
|
|
self.assertEqual((extension,media_type),('png','image/png'))
|
|
with Image.open(BytesIO(cleaned)) as checked: self.assertEqual(checked.mode,'RGBA')
|
|
|
|
def test_rejects_non_image_svg_and_oversized_input(self):
|
|
for contents in (b'<svg onload="alert(1)"></svg>', b'%PDF-1.7 fake document'):
|
|
with self.subTest(contents=contents[:4]), self.assertRaises(ScreenshotError):
|
|
sanitize_screenshot(contents)
|
|
with self.assertRaisesRegex(ScreenshotError,'too_large'):
|
|
sanitize_screenshot(b'X'*(MAX_SCREENSHOT_BYTES+1))
|
|
|
|
def test_rejects_excessive_dimensions(self):
|
|
source=self.png((5,5))
|
|
from unittest.mock import patch
|
|
with patch('safe_images.MAX_SCREENSHOT_PIXELS',16), self.assertRaisesRegex(ScreenshotError,'too_large'):
|
|
sanitize_screenshot(source)
|
|
|
|
|
|
if __name__ == '__main__': unittest.main()
|