ajyy/backend/app/core/upload_validate.py

"""上传文件校验：扩展名白名单与内容 magic。"""
from __future__ import annotations

import os
from typing import Callable, Optional

from app.core.response import AppHTTPError

# 对齐原型：图片、视频、常见办公文档与压缩包
ALLOWED_EXTENSIONS = {
    ".jpg",
    ".jpeg",
    ".png",
    ".gif",
    ".webp",
    ".bmp",
    ".mp4",
    ".mov",
    ".avi",
    ".webm",
    ".pdf",
    ".doc",
    ".docx",
    ".xls",
    ".xlsx",
    ".zip",
}

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}


def _starts_with(prefix: bytes) -> Callable[[bytes], bool]:
    def check(content: bytes) -> bool:
        return content.startswith(prefix)

    return check


def _starts_with_any(*prefixes: bytes) -> Callable[[bytes], bool]:
    def check(content: bytes) -> bool:
        return any(content.startswith(p) for p in prefixes)

    return check


def _webp(content: bytes) -> bool:
    return len(content) >= 12 and content[:4] == b"RIFF" and content[8:12] == b"WEBP"


_MAGIC_CHECKS: dict[str, Callable[[bytes], bool]] = {
    ".png": _starts_with(b"\x89PNG\r\n\x1a\n"),
    ".jpg": _starts_with(b"\xff\xd8\xff"),
    ".jpeg": _starts_with(b"\xff\xd8\xff"),
    ".gif": _starts_with_any(b"GIF87a", b"GIF89a"),
    ".webp": _webp,
    ".bmp": _starts_with(b"BM"),
    ".pdf": _starts_with(b"%PDF"),
}


def file_extension(filename: str) -> str:
    return os.path.splitext(filename or "")[1].lower()


def validate_upload(filename: Optional[str], content: bytes) -> None:
    if not filename:
        raise AppHTTPError(400, "请选择文件")
    ext = file_extension(filename)
    if ext not in ALLOWED_EXTENSIONS:
        raise AppHTTPError(400, "不支持的文件类型")
    magic_check = _MAGIC_CHECKS.get(ext)
    if magic_check is not None and not magic_check(content):
        raise AppHTTPError(400, "文件内容与扩展名不符")


def validate_image_upload(filename: Optional[str], content: bytes) -> None:
    if not filename:
        raise AppHTTPError(400, "请选择文件")
    ext = file_extension(filename)
    if ext not in IMAGE_EXTENSIONS:
        raise AppHTTPError(400, "不支持的文件类型")
    validate_upload(filename, content)