108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
import asyncio
|
|
import io
|
|
import logging
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from fastapi import HTTPException, UploadFile, status
|
|
from PIL import Image, UnidentifiedImageError
|
|
|
|
from app.config import settings
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
ALLOWED_FORMATS = {"JPEG": "jpg", "PNG": "png", "WEBP": "webp"}
|
|
|
|
|
|
def _process_image(raw: bytes) -> tuple[bytes, str, bytes]:
|
|
Image.MAX_IMAGE_PIXELS = settings.max_image_pixels
|
|
try:
|
|
with Image.open(io.BytesIO(raw)) as source:
|
|
if source.width * source.height > settings.max_image_pixels:
|
|
raise ValueError("图片像素尺寸过大")
|
|
source.load()
|
|
image_format = source.format
|
|
if image_format not in ALLOWED_FORMATS:
|
|
raise ValueError("仅支持 JPEG、PNG 和 WebP 图片")
|
|
|
|
extension = ALLOWED_FORMATS[image_format]
|
|
normalized = io.BytesIO()
|
|
if image_format == "PNG":
|
|
source.save(normalized, format="PNG", optimize=True)
|
|
else:
|
|
rgb = source.convert("RGB")
|
|
output_format = "JPEG" if image_format == "JPEG" else "WEBP"
|
|
rgb.save(normalized, format=output_format, quality=92, optimize=True)
|
|
|
|
thumbnail = source.copy()
|
|
thumbnail.thumbnail(
|
|
(settings.thumbnail_max_edge, settings.thumbnail_max_edge),
|
|
Image.Resampling.LANCZOS,
|
|
)
|
|
thumbnail_bytes = io.BytesIO()
|
|
thumbnail.convert("RGB").save(
|
|
thumbnail_bytes,
|
|
format="JPEG",
|
|
quality=78,
|
|
optimize=True,
|
|
)
|
|
return normalized.getvalue(), extension, thumbnail_bytes.getvalue()
|
|
except (UnidentifiedImageError, OSError, Image.DecompressionBombError, ValueError) as exc:
|
|
raise ValueError(str(exc) or "图片内容无效") from exc
|
|
|
|
|
|
def _write_files(
|
|
user_id: uuid.UUID,
|
|
original: bytes,
|
|
extension: str,
|
|
thumbnail: bytes,
|
|
) -> tuple[str, str]:
|
|
file_id = uuid.uuid4()
|
|
relative_dir = Path(str(user_id))
|
|
image_path = relative_dir / f"{file_id}.{extension}"
|
|
thumbnail_path = relative_dir / f"{file_id}-thumb.jpg"
|
|
target_dir = settings.upload_dir / relative_dir
|
|
target_dir.mkdir(parents=True, exist_ok=True)
|
|
(settings.upload_dir / image_path).write_bytes(original)
|
|
(settings.upload_dir / thumbnail_path).write_bytes(thumbnail)
|
|
return image_path.as_posix(), thumbnail_path.as_posix()
|
|
|
|
|
|
async def save_cloud_image(file: UploadFile, user_id: uuid.UUID) -> tuple[str, str]:
|
|
raw = await file.read(settings.max_upload_bytes + 1)
|
|
if not raw:
|
|
raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="图片不能为空")
|
|
if len(raw) > settings.max_upload_bytes:
|
|
raise HTTPException(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="图片文件过大")
|
|
try:
|
|
original, extension, thumbnail = await asyncio.to_thread(_process_image, raw)
|
|
except ValueError as exc:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
|
detail=f"图片无法处理：{exc}",
|
|
) from exc
|
|
return await asyncio.to_thread(_write_files, user_id, original, extension, thumbnail)
|
|
|
|
|
|
def _safe_file_path(relative_path: str) -> Path | None:
|
|
root = settings.upload_dir.resolve()
|
|
candidate = (root / relative_path).resolve()
|
|
if candidate == root or root not in candidate.parents:
|
|
return None
|
|
return candidate
|
|
|
|
|
|
async def delete_files(paths: list[str]) -> None:
|
|
def remove() -> None:
|
|
for relative_path in paths:
|
|
candidate = _safe_file_path(relative_path)
|
|
if not candidate:
|
|
logger.error("拒绝删除不安全的文件路径：%s", relative_path)
|
|
continue
|
|
try:
|
|
candidate.unlink(missing_ok=True)
|
|
except OSError:
|
|
logger.exception("删除图片失败：%s", relative_path)
|
|
|
|
await asyncio.to_thread(remove)
|