"""上传文件暂存、类型识别和安全落盘。""" 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