zip_processor.py

"""
zip_processor.py
------------------
Safe handling of user-uploaded ZIP archives containing multiple sprite
sheets.

Security measures:
  - Path traversal is blocked (entries like "../../etc/passwd" or absolute
    paths are rejected).
  - Nested ZIPs are ignored by default (not recursively extracted) --
    this avoids zip-bomb-via-nesting and matches the spec explicitly.
  - Per-file and total archive size are capped (configurable).
  - Only PNG/JPG/JPEG/WEBP are processed; SVG and anything else is skipped.
  - Files are extracted to memory / a scoped temp dir, never written using
    the archive's raw (untrusted) path.
"""

from __future__ import annotations

import zipfile
from dataclasses import dataclass, field
from pathlib import PurePosixPath

from config import (
    MAX_FILE_SIZE_MB,
    MAX_FILES_PER_ZIP,
    MAX_ZIP_SIZE_MB,
    REJECTED_EXTENSIONS,
    SUPPORTED_EXTENSIONS,
)


@dataclass
class ZipEntryResult:
    filename: str
    data: bytes


@dataclass
class ZipProcessingReport:
    entries: list[ZipEntryResult] = field(default_factory=list)
    skipped: list[str] = field(default_factory=list)  # (reason: filename)
    rejected_nested_zips: list[str] = field(default_factory=list)
    total_files_in_archive: int = 0


class ZipSecurityError(Exception):
    """Raised for hard-stop conditions (archive itself too large/malicious)."""


def _is_safe_member_path(name: str) -> bool:
    """
    Rejects absolute paths, drive letters, and any component that would
    escape the extraction root via '..'.
    """
    if not name or name.startswith("/") or name.startswith("\\"):
        return False
    # Windows drive letter, e.g. "C:\\"
    if len(name) > 1 and name[1] == ":":
        return False
    posix_path = PurePosixPath(name.replace("\\", "/"))
    if ".." in posix_path.parts:
        return False
    return True


def process_zip(zip_bytes: bytes, filename_hint: str = "upload.zip") -> ZipProcessingReport:
    archive_size_mb = len(zip_bytes) / (1024 * 1024)
    if archive_size_mb > MAX_ZIP_SIZE_MB:
        raise ZipSecurityError(
            f"ZIP '{filename_hint}' is {archive_size_mb:.1f}MB, exceeds MAX_ZIP_SIZE_MB={MAX_ZIP_SIZE_MB}MB"
        )

    report = ZipProcessingReport()

    import io

    try:
        zf = zipfile.ZipFile(io.BytesIO(zip_bytes))
    except zipfile.BadZipFile as exc:
        raise ZipSecurityError(f"'{filename_hint}' is not a valid ZIP file: {exc}") from exc

    infos = zf.infolist()
    report.total_files_in_archive = len(infos)

    if len(infos) > MAX_FILES_PER_ZIP:
        raise ZipSecurityError(
            f"ZIP contains {len(infos)} files, exceeds MAX_FILES_PER_ZIP={MAX_FILES_PER_ZIP}"
        )

    for info in infos:
        name = info.filename

        if info.is_dir():
            continue

        if not _is_safe_member_path(name):
            report.skipped.append(f"{name} (unsafe path, blocked)")
            continue

        # Guard against zip bombs: compare compressed vs uncompressed size.
        if info.file_size > MAX_FILE_SIZE_MB * 1024 * 1024:
            report.skipped.append(f"{name} (file too large, > {MAX_FILE_SIZE_MB}MB)")
            continue
        if info.compress_size > 0 and info.file_size / max(1, info.compress_size) > 200:
            # Suspiciously high compression ratio -- classic zip-bomb signature.
            report.skipped.append(f"{name} (suspicious compression ratio, blocked)")
            continue

        suffix = PurePosixPath(name).suffix.lower()

        if suffix in REJECTED_EXTENSIONS:
            report.skipped.append(f"{name} (SVG not supported)")
            continue

        if suffix == ".zip":
            report.rejected_nested_zips.append(name)
            continue

        if suffix not in SUPPORTED_EXTENSIONS:
            report.skipped.append(f"{name} (unsupported extension)")
            continue

        try:
            data = zf.read(info)
        except Exception as exc:  # noqa: BLE001
            report.skipped.append(f"{name} (could not read: {exc})")
            continue

        report.entries.append(ZipEntryResult(filename=PurePosixPath(name).name, data=data))

    return report