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

"""DMS存储目录创建和安全路径解析。"""

from __future__ import annotations

from pathlib import Path, PurePosixPath


STORAGE_SUBDIRECTORIES = (
    "original",
    "preview",
    "extracted",
    "temporary",
    "quarantine",
    "recycle",
)


class UnsafeStoragePathError(ValueError):
    """相对路径试图逃逸DMS存储根目录。"""


def ensure_storage_directories(storage_root: str | Path) -> Path:
    root = Path(storage_root).expanduser().resolve()
    root.mkdir(parents=True, exist_ok=True)
    for directory in STORAGE_SUBDIRECTORIES:
        (root / directory).mkdir(parents=True, exist_ok=True)
    return root


def resolve_storage_path(
    relative_path: str,
    storage_root: str | Path,
) -> Path:
    """安全解析数据库相对路径并保证结果仍位于存储根目录。"""
    if not relative_path or "\x00" in relative_path:
        raise UnsafeStoragePathError("存储相对路径不能为空或包含空字符")

    normalized = PurePosixPath(relative_path.replace("\\", "/"))
    has_drive_prefix = bool(normalized.parts and ":" in normalized.parts[0])
    if normalized.is_absolute() or has_drive_prefix or ".." in normalized.parts:
        raise UnsafeStoragePathError("存储路径必须是根目录内的安全相对路径")

    root = Path(storage_root).expanduser().resolve()
    candidate = root.joinpath(*normalized.parts).resolve()
    if not candidate.is_relative_to(root):
        raise UnsafeStoragePathError("存储路径超出DMS存储根目录")
    return candidate


def preview_cache_path(file_hash: str, storage_root: str | Path) -> Path:
    """基于文件SHA-256返回预览PDF缓存路径（支持前两位分层）。"""
    if not file_hash or len(file_hash) < 2:
        raise UnsafeStoragePathError("文件哈希无效")
    root = Path(storage_root).expanduser().resolve()
    prefix = file_hash[:2].lower()
    candidate = (root / "preview" / prefix / f"{file_hash}.pdf").resolve()
    if not candidate.is_relative_to(root / "preview"):
        raise UnsafeStoragePathError("预览缓存路径超出预览目录")
    return candidate


def preview_directory_for_hash(file_hash: str, storage_root: str | Path) -> Path:
    """返回预览缓存文件所在目录，确保位于 preview/ 下。"""
    path = preview_cache_path(file_hash, storage_root)
    return path.parent


def is_path_inside_preview(path: Path, storage_root: str | Path) -> bool:
    """检查路径是否位于预览缓存目录内。"""
    root = Path(storage_root).expanduser().resolve()
    resolved = path.resolve()
    return resolved.is_relative_to(root / "preview")