"""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()