document-management/backend/dms/storage/uploads.py

"""上传文件暂存、类型识别和安全落盘。"""

from __future__ import annotations

import hashlib
import uuid
import zipfile
from dataclasses import dataclass
from pathlib import Path, PurePosixPath

import olefile
from flask import current_app
from werkzeug.datastructures import FileStorage

from dms.common.errors import PayloadTooLargeError, UnsupportedFileTypeError
from dms.storage.envelope import promote_encrypt
from dms.storage.paths import ensure_storage_directories

ALLOWED_EXTENSIONS = {
    "doc", "docx", "pdf", "wps",
    "xls", "xlsx", "et",
    "jpg", "jpeg", "png", "gif", "bmp", "webp",
}
MIME_TYPES = {
    "doc": "application/msword",
    "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "pdf": "application/pdf",
    "wps": "application/vnd.ms-works",
    "xls": "application/vnd.ms-excel",
    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "et": "application/vnd.kingsoft.spreadsheet",
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "png": "image/png",
    "gif": "image/gif",
    "bmp": "image/bmp",
    "webp": "image/webp",
}
# 图片类扩展名：上传仅做魔数校验，预览时原图直出（不走 PDF 转换）
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "gif", "bmp", "webp"}
# 图片魔数签名：extension -> (签名, 起始偏移)；webp 为 RIFF 容器，偏移 8 处为 WEBP
_IMAGE_SIGNATURES: dict[str, tuple[bytes, int]] = {
    "jpg": (b"\xff\xd8\xff", 0),
    "jpeg": (b"\xff\xd8\xff", 0),
    "png": (b"\x89PNG\r\n\x1a\n", 0),
    "gif": (b"GIF8", 0),
    "bmp": (b"BM", 0),
    "webp": (b"WEBP", 8),
}
OLE_HEADER = bytes.fromhex("D0CF11E0A1B11AE1")


@dataclass(frozen=True, slots=True)
class StagedUpload:
    temporary_path: Path
    final_path: Path
    relative_path: str
    original_file_name: str
    extension: str
    mime_type: str
    file_size: int
    file_hash: str

    def cleanup(self) -> None:
        self.temporary_path.unlink(missing_ok=True)

    def cleanup_all(self) -> None:
        self.temporary_path.unlink(missing_ok=True)
        self.final_path.unlink(missing_ok=True)

    def promote(self) -> None:
        self.final_path.parent.mkdir(parents=True, exist_ok=True)
        # 契约1.21：original/ 且已配置 KEK 时按信封格式加密落盘，否则明文迁移
        if promote_encrypt(
            self.temporary_path, self.final_path, self.relative_path
        ):
            return
        self.temporary_path.replace(self.final_path)


def _safe_zip(path: Path, required: set[str]) -> None:
    try:
        with zipfile.ZipFile(path) as archive:
            if len(archive.infolist()) > 10000:
                raise UnsupportedFileTypeError("Office容器条目数量异常")
            uncompressed_size = sum(item.file_size for item in archive.infolist())
            if uncompressed_size > current_app.config["DMS_MAX_FILE_SIZE"] * 20:
                raise UnsupportedFileTypeError("Office容器解压规模异常")
            names = set(archive.namelist())
            for name in names:
                normalized = PurePosixPath(name.replace("\\", "/"))
                if normalized.is_absolute() or ".." in normalized.parts:
                    raise UnsupportedFileTypeError("Office容器包含不安全路径")
            if not required <= names:
                raise UnsupportedFileTypeError("Office容器结构与扩展名不一致")
            if archive.testzip() is not None:
                raise UnsupportedFileTypeError("Office ZIP容器损坏")
    except (zipfile.BadZipFile, OSError, RuntimeError) as exc:
        raise UnsupportedFileTypeError("Office ZIP容器损坏") from exc


def _validate_image(prefix: bytes, extension: str) -> None:
    signature, offset = _IMAGE_SIGNATURES[extension]
    valid = prefix[offset : offset + len(signature)] == signature
    if extension == "webp":
        valid = valid and prefix.startswith(b"RIFF")
    if not valid:
        raise UnsupportedFileTypeError("图片文件头与扩展名不一致")


def _validate(path: Path, extension: str) -> None:
    # 读取 12 字节：webp 需校验偏移 8~11 处的 WEBP 标识，其余类型不受影响
    prefix = path.read_bytes()[:12]
    if extension == "pdf":
        if not prefix.startswith(b"%PDF-"):
            raise UnsupportedFileTypeError("PDF文件头与扩展名不一致")
        return
    if extension == "docx":
        _safe_zip(path, {"[Content_Types].xml", "word/document.xml"})
        return
    if extension == "xlsx":
        _safe_zip(path, {"[Content_Types].xml", "xl/workbook.xml"})
        return
    if extension in IMAGE_EXTENSIONS:
        _validate_image(prefix, extension)
        return
    if not prefix.startswith(OLE_HEADER):
        raise UnsupportedFileTypeError("OLE文件头与扩展名不一致")
    try:
        with olefile.OleFileIO(path) as container:
            streams = {"/".join(parts) for parts in container.listdir()}
    except (OSError, IOError) as exc:
        raise UnsupportedFileTypeError("OLE/CFB容器损坏") from exc
    if extension == "doc" and "WordDocument" not in streams:
        raise UnsupportedFileTypeError("OLE容器不是有效DOC")
    if extension == "xls" and "Workbook" not in streams:
        raise UnsupportedFileTypeError("OLE容器不是有效XLS")
    # wps/et：与 DOC/XLS 同族的 OLE 容器，但内部流名不保证 WordDocument/Workbook，
    # 仅校验容器有效性；内容合法性由 LibreOffice 转换环节兜底（转换失败返回预览不可用）。


def stage_upload(file: FileStorage) -> StagedUpload:
    original_name = Path(file.filename or "").name
    if not original_name or "." not in original_name:
        raise UnsupportedFileTypeError("文件名缺少允许的扩展名")
    extension = original_name.rsplit(".", 1)[1].lower()
    if extension not in ALLOWED_EXTENSIONS:
        raise UnsupportedFileTypeError()
    root = ensure_storage_directories(current_app.config["DMS_STORAGE_ROOT"])
    temporary = root / "temporary" / f"{uuid.uuid4()}.upload"
    final_name = f"{uuid.uuid4()}.{extension}"
    final = root / "original" / final_name
    digest = hashlib.sha256()
    size = 0
    try:
        with temporary.open("xb") as stream:
            while chunk := file.stream.read(1024 * 1024):
                size += len(chunk)
                if size > current_app.config["DMS_MAX_FILE_SIZE"]:
                    raise PayloadTooLargeError("单个文件超过允许大小")
                digest.update(chunk)
                stream.write(chunk)
        if size == 0:
            raise UnsupportedFileTypeError("不允许上传空文件")
        _validate(temporary, extension)
        return StagedUpload(
            temporary_path=temporary,
            final_path=final,
            relative_path=f"original/{final_name}",
            original_file_name=original_name,
            extension=extension,
            mime_type=MIME_TYPES[extension],
            file_size=size,
            file_hash=digest.hexdigest(),
        )
    except Exception:
        temporary.unlink(missing_ok=True)
        raise