uploads.py

"""安全复制 Blender 项目及资源 ZIP；任何上传内容都不会被执行。"""

from __future__ import annotations

import os
import shutil
import stat
import tempfile
import zipfile
from pathlib import Path
from typing import BinaryIO

from render_config import positive_env_int, safe_relative_path


_CHUNK_SIZE = 1024 * 1024
_SCRIPT_SUFFIXES = {".py", ".pyc", ".pyo", ".pyw"}
_TAR_SUFFIXES = (".tar", ".tar.gz", ".tar.bz2", ".tar.xz", ".tgz", ".tbz", ".tbz2", ".txz")


class UploadError(ValueError):
    """可直接向调用者展示的上传错误。"""


class _Staging:
    def __init__(self, root: Path, byte_limit: int, file_limit: int):
        self.root = root
        self.byte_limit = byte_limit
        self.file_limit = file_limit
        self.bytes_written = 0
        self.entry_count = 0
        # 文件和文件夹共用名称空间，大小写不同也视为冲突。
        self.entries: dict[str, tuple[str, bool, bool]] = {}
        self.blends: list[str] = []

    def register(self, name: str, directory: bool = False) -> Path:
        self.entry_count += 1
        if self.entry_count > self.file_limit:
            raise UploadError("上传条目数超过服务端 MAX_ARCHIVE_FILES 限制。")
        try:
            safe_relative_path(name)
        except ValueError as exc:
            raise UploadError(f"不安全的资源路径 {name!r}：{exc}") from exc
        if not directory and Path(name).suffix.casefold() in _SCRIPT_SUFFIXES:
            raise UploadError("不允许上传或解压 Python 脚本文件。")
        if not directory and name.casefold().endswith(_TAR_SUFFIXES):
            raise UploadError("不支持 TAR 压缩包，请使用 ZIP。")
        components = name.split("/")
        for index in range(1, len(components)):
            parent = "/".join(components[:index])
            folded = parent.casefold()
            existing = self.entries.get(folded)
            if existing and (existing[0] != parent or not existing[1]):
                raise UploadError(f"路径存在大小写或文件/文件夹冲突：{parent}")
            if not existing:
                self.entries[folded] = (parent, True, False)
        folded = name.casefold()
        existing = self.entries.get(folded)
        # 已创建的隐式父文件夹允许一次显式 ZIP 目录条目。
        if existing and not (
            directory and existing == (name, True, False)
        ):
            raise UploadError(f"上传包含重复或冲突路径：{name}")
        self.entries[folded] = (name, directory, True)
        target = self.root.joinpath(*components)
        if directory:
            target.mkdir(parents=True, exist_ok=True)
        else:
            target.parent.mkdir(parents=True, exist_ok=True)
            if name.casefold().endswith(".blend"):
                self.blends.append(name)
        return target

    def copy_stream(self, source: BinaryIO, target: Path, declared_size: int) -> None:
        if declared_size < 0 or self.bytes_written + declared_size > self.byte_limit:
            raise UploadError("解压/暂存总大小超过服务端 MAX_EXTRACT_MB 限制。")
        actual = 0
        with target.open("xb") as output:
            while chunk := source.read(_CHUNK_SIZE):
                actual += len(chunk)
                self.bytes_written += len(chunk)
                if self.bytes_written > self.byte_limit or actual > declared_size:
                    raise UploadError("文件实际大小超过声明值或服务端大小限制。")
                output.write(chunk)
        if actual != declared_size:
            raise UploadError("上传文件不完整，实际大小与声明值不一致。")

    def unpack(self, source: BinaryIO) -> None:
        with zipfile.ZipFile(source) as archive:
            infos = archive.infolist()
            if len(infos) + self.entry_count > self.file_limit:
                raise UploadError("压缩包条目数超过服务端 MAX_ARCHIVE_FILES 限制。")
            for info in infos:
                # ZipInfo.filename 会截断 NUL，必须检查保存的原始名称。
                original = getattr(info, "orig_filename", info.filename)
                if original != info.filename or "\x00" in original:
                    raise UploadError("ZIP 文件名包含 NUL 字符。")
                mode = (info.external_attr >> 16) & 0xFFFF
                kind = stat.S_IFMT(mode)
                if kind not in (0, stat.S_IFREG, stat.S_IFDIR):
                    raise UploadError("ZIP 中不允许符号链接或特殊文件。")
                directory = info.is_dir()
                if kind == stat.S_IFDIR and not directory:
                    raise UploadError("ZIP 文件夹属性与路径不一致。")
                if info.flag_bits & 1:
                    raise UploadError("不支持加密 ZIP。")
                if info.file_size / max(info.compress_size, 1) > 500:
                    raise UploadError("ZIP 压缩比超过 500，拒绝潜在压缩炸弹。")
                name = original[:-1] if directory else original
                target = self.register(name, directory)
                if directory:
                    if info.file_size:
                        raise UploadError("ZIP 文件夹条目不能包含文件数据。")
                    continue
                # 内嵌 ZIP 作为普通资源复制，不递归解压。
                with archive.open(info, "r") as member:
                    self.copy_stream(member, target, info.file_size)


def _open_regular_file(source: Path) -> BinaryIO:
    source_mode = source.lstat().st_mode
    if stat.S_ISLNK(source_mode):
        raise UploadError("不允许上传符号链接。")
    if not stat.S_ISREG(source_mode):
        raise UploadError("上传对象必须是现有普通文件。")
    # NONBLOCK 避免检查后被替换为 FIFO 的路径阻塞任务。
    flags = (os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
             | getattr(os, "O_NONBLOCK", 0) | getattr(os, "O_BINARY", 0))
    descriptor = os.open(source, flags)
    try:
        if not stat.S_ISREG(os.fstat(descriptor).st_mode):
            raise UploadError("上传对象必须是现有普通文件。")
        return os.fdopen(descriptor, "rb")
    except BaseException:
        os.close(descriptor)
        raise


def stage_uploads(files: list[str], project_dir: Path, blend_file: str | None) -> Path:
    """原子暂存项目，成功返回所选 .blend 的绝对路径。

    project_dir 必须不存在或为空。失败仅清理本次创建的临时目录，
    不会删除调用者已有文件；项目目录不会出现部分解压结果。
    """
    project_dir = Path(project_dir).absolute()
    temporary: Path | None = None
    try:
        file_limit = positive_env_int("MAX_ARCHIVE_FILES", 5000)
        upload_limit = positive_env_int("MAX_UPLOAD_MB", 1024) * 1024 * 1024
        extract_limit = positive_env_int("MAX_EXTRACT_MB", 4096) * 1024 * 1024
        if not files:
            raise UploadError("请上传 .blend 文件及所需资源，或包含项目的 ZIP。")
        if len(files) > file_limit:
            raise UploadError("上传文件数量超过服务端 MAX_ARCHIVE_FILES 限制。")
        if blend_file is not None:
            safe_relative_path(blend_file)
            if not blend_file.casefold().endswith(".blend"):
                raise UploadError("blend_file 必须指定 .blend 文件。")
        if project_dir.is_symlink():
            raise UploadError("项目目录不能是符号链接。")
        if project_dir.exists() and (
            not project_dir.is_dir() or any(project_dir.iterdir())
        ):
            raise UploadError("项目目录必须为空，不能覆盖已有文件。")
        project_dir.parent.mkdir(parents=True, exist_ok=True)
        temporary = Path(tempfile.mkdtemp(prefix=f".{project_dir.name}.upload-", dir=project_dir.parent))
        staging = _Staging(temporary, extract_limit, file_limit)
        uploaded_bytes = 0
        upload_names: set[str] = set()
        for file_path in files:
            source = Path(file_path)
            name = source.name
            safe_relative_path(name)
            if name.casefold() in upload_names:
                raise UploadError(f"上传文件名称重复：{name}")
            upload_names.add(name.casefold())
            if source.suffix.casefold() in _SCRIPT_SUFFIXES:
                raise UploadError("不允许上传 Python 脚本文件。")
            if name.casefold().endswith(_TAR_SUFFIXES):
                raise UploadError("不支持 TAR 压缩包，请使用 ZIP。")
            with _open_regular_file(source) as handle:
                size = os.fstat(handle.fileno()).st_size
                uploaded_bytes += size
                if uploaded_bytes > upload_limit:
                    raise UploadError("上传总大小超过服务端 MAX_UPLOAD_MB 限制。")
                if source.suffix.casefold() == ".zip":
                    staging.unpack(handle)
                else:
                    staging.copy_stream(handle, staging.register(name), size)
        if not staging.blends:
            raise UploadError("上传内容中找不到 .blend 文件。")
        if blend_file is None:
            if len(staging.blends) != 1:
                raise UploadError("存在多个 .blend 文件，请通过 blend_file 指定相对路径。")
            selected = staging.blends[0]
        else:
            if blend_file not in staging.blends:
                raise UploadError(f"找不到 blend_file 指定的文件：{blend_file}")
            selected = blend_file
        if project_dir.exists():
            project_dir.rmdir()  # 只移除空目录；发生竞争写入时安全失败。
        os.replace(temporary, project_dir)
        temporary = None
        return project_dir.joinpath(*selected.split("/"))
    except UploadError:
        raise
    except (OSError, ValueError, RuntimeError, zipfile.BadZipFile, NotImplementedError) as exc:
        raise UploadError(f"无法暂存上传项目：{exc}") from exc
    finally:
        if temporary is not None:
            shutil.rmtree(temporary, ignore_errors=True)