#!/usr/bin/env python3
"""Shared, standard-library-only ZIP inspection and safe extraction.

This module is intentionally safe to ship inside public online-skill packages.
Both the release-side validator and the installed update client must use this
same implementation so an archive cannot pass one trust boundary and bypass
the other.
"""

from __future__ import annotations

import stat
import unicodedata
import zipfile
from pathlib import Path, PurePosixPath
from typing import Iterable


MAX_ARCHIVE_BYTES = 256 * 1024 * 1024
MAX_MEMBER_COUNT = 5_000
MAX_MEMBER_NAME_BYTES = 1_024
MAX_MEMBER_UNCOMPRESSED_BYTES = 32 * 1024 * 1024
MAX_TOTAL_UNCOMPRESSED_BYTES = 128 * 1024 * 1024
MAX_COMPRESSION_RATIO = 200.0
ALLOWED_COMPRESSION_METHODS = {zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED}


class ArchiveSafetyError(ValueError):
    """Raised when a ZIP is ambiguous, unsafe, or too expensive to extract."""


def _unsafe_member_path(name: str) -> bool:
    if not name or "\x00" in name or "\\" in name or name.startswith(("/", "~")):
        return True
    path = PurePosixPath(name)
    if path.is_absolute():
        return True
    parts = name.split("/")
    return any(part in {"", ".", ".."} for part in parts[:-1]) or any(
        part in {".", ".."} for part in parts
    )


def _normalised_member_name(name: str) -> tuple[str, str]:
    if not isinstance(name, str) or _unsafe_member_path(name):
        raise ArchiveSafetyError(f"ZIP contains an unsafe member path: {name!r}")
    normalised = unicodedata.normalize("NFC", name)
    if normalised != name:
        raise ArchiveSafetyError(
            f"ZIP member name is not Unicode NFC-normalized: {name!r}"
        )
    if len(normalised.encode("utf-8")) > MAX_MEMBER_NAME_BYTES:
        raise ArchiveSafetyError("ZIP member name exceeds the safety limit")
    return normalised, normalised.casefold()


def _member_type(member: zipfile.ZipInfo) -> int:
    if member.create_system != 3:
        return 0
    return stat.S_IFMT(member.external_attr >> 16)


def inspect_archive_path(path: Path) -> Path:
    """Validate the archive filesystem object and return its resolved path."""

    expanded = path.expanduser()
    if expanded.is_symlink():
        raise ArchiveSafetyError(f"ZIP path must not be a symlink: {path}")
    try:
        resolved = expanded.resolve(strict=True)
    except OSError as exc:
        raise ArchiveSafetyError(f"ZIP is missing: {path}") from exc
    if not resolved.is_file():
        raise ArchiveSafetyError(f"ZIP is not a regular file: {path}")
    if resolved.stat().st_size > MAX_ARCHIVE_BYTES:
        raise ArchiveSafetyError("ZIP exceeds the compressed size limit")
    return resolved


def inspect_archive_members(
    archive: zipfile.ZipFile,
) -> list[tuple[zipfile.ZipInfo, str]]:
    """Fail closed on ambiguous names, special entries, and ZIP bombs."""

    members = archive.infolist()
    if len(members) > MAX_MEMBER_COUNT:
        raise ArchiveSafetyError(
            f"ZIP has too many members: {len(members)} > {MAX_MEMBER_COUNT}"
        )
    seen: dict[str, str] = {}
    total = 0
    inspected: list[tuple[zipfile.ZipInfo, str]] = []
    for member in members:
        normalised, collision_key = _normalised_member_name(member.filename)
        previous = seen.get(collision_key)
        if previous is not None:
            raise ArchiveSafetyError(
                "ZIP contains duplicate/casefold/Unicode-colliding members: "
                f"{previous!r} and {member.filename!r}"
            )
        seen[collision_key] = member.filename
        if member.flag_bits & 0x1:
            raise ArchiveSafetyError(
                f"ZIP contains an encrypted member: {member.filename}"
            )
        if member.compress_type not in ALLOWED_COMPRESSION_METHODS:
            raise ArchiveSafetyError(
                f"ZIP uses an unsupported compression method: {member.filename}"
            )
        member_type = _member_type(member)
        is_directory = member.is_dir() or normalised.endswith("/")
        if member_type == stat.S_IFLNK:
            raise ArchiveSafetyError(f"ZIP contains a symlink: {member.filename}")
        if member_type not in {0, stat.S_IFREG, stat.S_IFDIR}:
            raise ArchiveSafetyError(
                f"ZIP contains a special filesystem entry: {member.filename}"
            )
        if is_directory and member_type == stat.S_IFREG:
            raise ArchiveSafetyError(
                f"ZIP directory has a conflicting file mode: {member.filename}"
            )
        if not is_directory and member_type == stat.S_IFDIR:
            raise ArchiveSafetyError(
                f"ZIP file has a conflicting directory mode: {member.filename}"
            )
        if member.file_size < 0 or member.compress_size < 0:
            raise ArchiveSafetyError("ZIP contains an invalid member size")
        if member.file_size > MAX_MEMBER_UNCOMPRESSED_BYTES:
            raise ArchiveSafetyError(
                f"ZIP member exceeds the uncompressed size limit: {member.filename}"
            )
        total += member.file_size
        if total > MAX_TOTAL_UNCOMPRESSED_BYTES:
            raise ArchiveSafetyError("ZIP exceeds the total uncompressed size limit")
        if member.file_size:
            if member.compress_size == 0:
                raise ArchiveSafetyError(
                    f"ZIP member has an impossible compression size: {member.filename}"
                )
            ratio = member.file_size / member.compress_size
            if ratio > MAX_COMPRESSION_RATIO:
                raise ArchiveSafetyError(
                    "ZIP member exceeds the compression-ratio limit: "
                    f"{member.filename} ({ratio:.1f}:1)"
                )
        inspected.append((member, normalised))
    return inspected


def safe_extract_archive(
    archive: zipfile.ZipFile,
    members: Iterable[tuple[zipfile.ZipInfo, str]],
    destination: Path,
) -> None:
    """Extract inspected members without ``extractall`` or overwrite behavior."""

    expanded = destination.expanduser()
    if expanded.is_symlink():
        raise ArchiveSafetyError("ZIP extraction root must not be a symlink")
    expanded.mkdir(parents=True, exist_ok=True)
    destination = expanded.resolve(strict=True)
    for member, normalised in members:
        relative = PurePosixPath(normalised)
        target = destination.joinpath(*relative.parts)
        resolved_parent = target.parent.resolve(strict=False)
        if not resolved_parent.is_relative_to(destination):
            raise ArchiveSafetyError(
                f"ZIP extraction path escapes the target root: {normalised}"
            )
        if member.is_dir() or normalised.endswith("/"):
            if target.exists() and not target.is_dir():
                raise ArchiveSafetyError(
                    f"ZIP directory collides with an existing file: {normalised}"
                )
            target.mkdir(parents=True, exist_ok=True)
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        written = 0
        try:
            with archive.open(member, "r") as source, target.open("xb") as output:
                while True:
                    chunk = source.read(1024 * 1024)
                    if not chunk:
                        break
                    written += len(chunk)
                    if written > member.file_size:
                        raise ArchiveSafetyError(
                            "ZIP member expanded beyond its declared size: "
                            f"{normalised}"
                        )
                    output.write(chunk)
        except FileExistsError as exc:
            raise ArchiveSafetyError(
                f"ZIP extraction would overwrite an existing path: {normalised}"
            ) from exc
        if written != member.file_size:
            raise ArchiveSafetyError(
                f"ZIP member size changed during extraction: {normalised}"
            )
