Spaces:
Running
Running
File size: 4,168 Bytes
39ff632 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """Upload validation/normalisation tests (app/image_guard.py)."""
from __future__ import annotations
import io
import struct
import zlib
import pytest
from PIL import Image, ImageFile
from app.image_guard import ImageRejected, validate_and_normalize
def make_png_bytes(w: int = 64, h: int = 64, color=(200, 30, 30)) -> bytes:
raw = bytearray()
for y in range(h):
raw.append(0)
for _ in range(w):
raw += bytes(color)
def chunk(tag, data):
c = struct.pack(">I", len(data)) + tag + data
return c + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr)
+ chunk(b"IDAT", zlib.compress(bytes(raw), 9)) + chunk(b"IEND", b""))
class TestAccepts:
def test_png_accepted_and_normalized(self):
result = validate_and_normalize(make_png_bytes())
assert result.original_format == "PNG"
assert result.width == 64 and result.height == 64
# Output really is a fresh RGB PNG.
img = Image.open(io.BytesIO(result.png_bytes))
assert img.format == "PNG" and img.mode == "RGB"
def test_jpeg_accepted(self):
buf = io.BytesIO()
Image.new("RGB", (50, 40), (10, 120, 200)).save(buf, "JPEG")
result = validate_and_normalize(buf.getvalue())
assert result.original_format == "JPEG"
def test_gif_accepted_first_frame(self):
buf = io.BytesIO()
Image.new("P", (32, 32)).save(buf, "GIF")
result = validate_and_normalize(buf.getvalue())
assert result.original_format == "GIF"
def test_downscale_applies(self):
result = validate_and_normalize(make_png_bytes(2000, 1000), normalize_max_side=1024)
assert result.downscaled is True
assert max(result.width, result.height) <= 1024
def test_exif_stripped(self):
img = Image.new("RGB", (24, 24), (1, 2, 3))
buf = io.BytesIO()
img.save(buf, "PNG", pnginfo=None)
result = validate_and_normalize(buf.getvalue())
assert b"Exif" not in result.png_bytes
class TestRejects:
def test_empty(self):
with pytest.raises(ImageRejected) as exc:
validate_and_normalize(b"")
assert exc.value.code == "empty_file"
def test_html_bytes(self):
with pytest.raises(ImageRejected) as exc:
validate_and_normalize(b"<html><body>nope</body></html>")
assert exc.value.code in {"svg_rejected", "not_an_image"}
def test_svg_rejected(self):
svg = b'<?xml version="1.0"?><svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>'
with pytest.raises(ImageRejected) as exc:
validate_and_normalize(svg)
assert exc.value.code == "svg_rejected"
def test_oversize_bytes(self):
with pytest.raises(ImageRejected) as exc:
validate_and_normalize(b"\x89PNG" + b"0" * 2048, max_bytes=1024)
assert exc.value.code == "file_too_large"
def test_truncated_png(self):
with pytest.raises(ImageRejected):
validate_and_normalize(make_png_bytes()[:60])
def test_pixel_cap(self):
# Declared dimensions exceed the cap; Pillow raises on open.
with pytest.raises(ImageRejected):
validate_and_normalize(make_png_bytes(64, 64), max_pixels=100)
def test_dimension_cap_precedes_pixel_decode(self, monkeypatch):
decoded = False
original_load = ImageFile.ImageFile.load
def tracked_load(image, *args, **kwargs):
nonlocal decoded
decoded = True
return original_load(image, *args, **kwargs)
monkeypatch.setattr(ImageFile.ImageFile, "load", tracked_load)
with pytest.raises(ImageRejected) as exc:
validate_and_normalize(
make_png_bytes(64, 64), max_pixels=1_000_000, max_side=32
)
assert exc.value.code == "image_too_large"
assert decoded is False
def test_random_bytes(self):
with pytest.raises(ImageRejected):
validate_and_normalize(b"\xde\xad\xbe\xef" * 64)
|