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.paths import ensure_storage_directories

ALLOWED_EXTENSIONS = {"doc", "docx", "pdf", "xls", "xlsx"}
MIME_TYPES = {
    "doc": "application/msword",
    "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "pdf": "application/pdf",
    "xls": "application/vnd.ms-excel",
    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
}
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)
        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(path: Path, extension: str) -> None:
    prefix = path.read_bytes()[:8]
    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 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 not ({"Workbook", "Book"} & streams):
        raise UnsupportedFileTypeError("OLE容器不是有效XLS")


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