webapi/image_guard.py

"""
webapi/image_guard.py — ตรวจรูปที่ผู้ใช้อัปโหลดก่อนถึงบอท

  · จำกัดขนาดไฟล์ (ตรวจระหว่างอ่าน ไม่อ่านทั้งไฟล์ก่อนค่อยตรวจ)
  · ยอมรับเฉพาะ JPEG / PNG / WEBP ดูจากเนื้อไฟล์ ไม่เชื่อนามสกุลหรือ content-type
  · กัน decompression bomb (จำนวนพิกเซล)
  · หมุนตาม EXIF · ย่อด้านยาวสุดไม่เกิน max_side · แปลงเป็น RGB (เหมือน run_testcases.py)
"""
from __future__ import annotations

import io
from dataclasses import dataclass
from typing import Optional

from PIL import Image, ImageOps, UnidentifiedImageError

ALLOWED_FORMATS = {"JPEG", "PNG", "WEBP"}


class ImageRejected(ValueError):
    def __init__(self, code: str, message_th: str):
        super().__init__(message_th)
        self.code = code
        self.message_th = message_th


@dataclass
class CleanImage:
    image: Image.Image
    original_format: str
    original_size: tuple[int, int]
    bytes_in: int


def load_image(data: bytes, max_bytes: int, max_pixels: int, max_side: int) -> CleanImage:
    if len(data) > max_bytes:
        raise ImageRejected("too_large", f"รูปใหญ่เกิน {max_bytes // (1024 * 1024)} MB ครับ รบกวนย่อรูปแล้วส่งใหม่")
    if not data:
        raise ImageRejected("empty", "ไฟล์รูปว่างครับ")
    try:
        probe = Image.open(io.BytesIO(data))
        fmt = probe.format or ""
        w, h = probe.size
    except (UnidentifiedImageError, OSError):
        raise ImageRejected("not_image", "ไฟล์นี้ไม่ใช่รูปภาพที่ระบบอ่านได้ครับ รองรับ JPG, PNG, WEBP")
    if fmt not in ALLOWED_FORMATS:
        raise ImageRejected("format", f"ยังไม่รองรับรูปแบบ {fmt or 'นี้'} ครับ รองรับ JPG, PNG, WEBP")
    if w * h > max_pixels:
        raise ImageRejected("too_many_pixels", "รูปมีความละเอียดสูงเกินไปครับ รบกวนย่อรูปแล้วส่งใหม่")
    try:
        img = Image.open(io.BytesIO(data))
        img.load()
    except Exception:
        raise ImageRejected("corrupt", "ไฟล์รูปเสียหรืออ่านไม่ครบครับ รบกวนส่งใหม่")
    img = ImageOps.exif_transpose(img)
    img = img.convert("RGB")
    if max(img.size) > max_side:
        img.thumbnail((max_side, max_side), Image.LANCZOS)
    return CleanImage(img, fmt, (w, h), len(data))


async def read_limited(upload, max_bytes: int) -> bytes:
    """อ่าน UploadFile ทีละก้อน หยุดทันทีเมื่อเกินขนาด"""
    buf = bytearray()
    while True:
        chunk = await upload.read(64 * 1024)
        if not chunk:
            break
        buf += chunk
        if len(buf) > max_bytes:
            raise ImageRejected("too_large", f"รูปใหญ่เกิน {max_bytes // (1024 * 1024)} MB ครับ รบกวนย่อรูปแล้วส่งใหม่")
    return bytes(buf)