66 lines
2.6 KiB
Python
66 lines
2.6 KiB
Python
"""Decode and re-encode small raster uploads before they leave MetalCircle."""
|
|
from io import BytesIO
|
|
import warnings
|
|
|
|
from PIL import Image, ImageOps
|
|
|
|
|
|
MAX_SCREENSHOT_BYTES = 4 * 1024 * 1024
|
|
MAX_SCREENSHOT_PIXELS = 16_000_000
|
|
MAX_SCREENSHOT_OUTPUT_BYTES = 8 * 1024 * 1024
|
|
|
|
|
|
class ScreenshotError(ValueError):
|
|
"""Stable validation code; never includes the uploaded filename or bytes."""
|
|
def __init__(self, code):
|
|
self.code = code
|
|
super().__init__(code)
|
|
|
|
|
|
def sanitize_screenshot(contents):
|
|
if not contents:
|
|
raise ScreenshotError('empty')
|
|
if len(contents) > MAX_SCREENSHOT_BYTES:
|
|
raise ScreenshotError('too_large')
|
|
|
|
try:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter('error', Image.DecompressionBombWarning)
|
|
with Image.open(BytesIO(contents)) as candidate:
|
|
if candidate.format not in {'JPEG', 'PNG', 'WEBP'}:
|
|
raise ScreenshotError('unsupported')
|
|
if candidate.width * candidate.height > MAX_SCREENSHOT_PIXELS:
|
|
raise ScreenshotError('too_large')
|
|
if getattr(candidate, 'n_frames', 1) != 1:
|
|
raise ScreenshotError('unsupported')
|
|
candidate.verify()
|
|
with Image.open(BytesIO(contents)) as candidate:
|
|
if candidate.width * candidate.height > MAX_SCREENSHOT_PIXELS:
|
|
raise ScreenshotError('too_large')
|
|
image = ImageOps.exif_transpose(candidate)
|
|
image.load()
|
|
has_alpha = image.mode in {'RGBA', 'LA'} or (
|
|
image.mode == 'P' and 'transparency' in image.info
|
|
)
|
|
if image.mode not in {'RGB', 'RGBA'}:
|
|
image = image.convert('RGBA' if has_alpha else 'RGB')
|
|
image.thumbnail((1600, 1600))
|
|
output = BytesIO()
|
|
if image.mode == 'RGBA':
|
|
image.save(output, format='PNG', optimize=True)
|
|
extension, media_type = 'png', 'image/png'
|
|
else:
|
|
image.save(output, format='JPEG', quality=85, optimize=True)
|
|
extension, media_type = 'jpg', 'image/jpeg'
|
|
safe_contents = output.getvalue()
|
|
except ScreenshotError:
|
|
raise
|
|
except Exception:
|
|
# Decoders may raise different exception types for malformed inputs.
|
|
# Fail closed without exposing decoder details to the caller.
|
|
raise ScreenshotError('invalid') from None
|
|
|
|
if len(safe_contents) > MAX_SCREENSHOT_OUTPUT_BYTES:
|
|
raise ScreenshotError('too_large')
|
|
return safe_contents, extension, media_type
|