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
|
"""Safe conversion between MCP Base64 inputs and Pillow images."""
from __future__ import annotations
import base64
import binascii
import io
import json
from typing import Any
from PIL import Image, PngImagePlugin, UnidentifiedImageError
from .errors import InvalidImageError
MAX_ENCODED_LENGTH = 36_000_000
MAX_PIXELS = 40_000_000
SUPPORTED_FORMATS = frozenset({"PNG", "JPEG"})
def decode_image(value: str, *, mask: bool = False) -> Image.Image:
"""Decode a bounded PNG/JPEG Base64 value or data URI and fully validate it."""
if not isinstance(value, str) or len(value) > MAX_ENCODED_LENGTH:
raise InvalidImageError("image Base64 input is missing or too large")
payload = value.strip()
if payload.startswith("data:"):
try:
header, payload = payload.split(",", 1)
except ValueError as exc:
raise InvalidImageError("malformed image data URI") from exc
if header.lower() not in {
"data:image/png;base64",
"data:image/jpeg;base64",
"data:image/jpg;base64",
}:
raise InvalidImageError("data URI must contain a Base64 PNG or JPEG")
try:
raw = base64.b64decode(payload, validate=True)
except (binascii.Error, ValueError) as exc:
raise InvalidImageError("image is not valid Base64") from exc
try:
with Image.open(io.BytesIO(raw)) as opened:
if opened.format not in SUPPORTED_FORMATS:
raise InvalidImageError("only PNG and JPEG images are supported")
if opened.width * opened.height > MAX_PIXELS:
raise InvalidImageError("image pixel count exceeds the safety limit")
opened.load()
return opened.convert("L" if mask else "RGB")
except InvalidImageError:
raise
except (OSError, UnidentifiedImageError, Image.DecompressionBombError) as exc:
raise InvalidImageError("input is not a valid PNG or JPEG image") from exc
def encode_png(image: Image.Image, metadata: dict[str, Any]) -> bytes:
"""Serialize one RGB image as PNG with stable textual generation metadata."""
info = PngImagePlugin.PngInfo()
for key, value in metadata.items():
text = json.dumps(value) if isinstance(value, (dict, list)) else str(value)
info.add_text(str(key), text)
output = io.BytesIO()
image.convert("RGB").save(output, format="PNG", pnginfo=info)
return output.getvalue()
|
