Blame for src/nida/data/remote_zip.py

"""List and extract single members of a *remote* zip archive with HTTP ``Range`` requests.

EARS ships one ~564 MB zip per speaker (ADR-0007) of which only a small part is needed. This
module reads the end-of-central-directory record (classic and ZIP64), the central directory
and then only the byte ranges of the requested members, so the archive is never downloaded
as a whole.

* :class:`RemoteZip` opens a URL. :meth:`RemoteZip.members` lists the archive,
  :meth:`RemoteZip.extract` writes one member to a file (atomically: temporary file in the
  destination folder, fsync, rename; a partial member never looks complete).
* Supported: stored (0) and deflate (8) members, ZIP64 end records and ZIP64 extra fields,
  UTF-8 and CP437 names. Not supported (clear :class:`RemoteZipError`): encryption, other
  compression methods, multi-disk archives, archives with data prepended to the zip.
* Every read goes through one *read-ahead window* (default 4 MiB): members that are close to
  each other in the archive are served by a single request. Extract in ascending
  ``header_offset`` order to profit from it.
* Network access goes through the injectable :data:`nida.bench.fetch.Opener`; transient
  failures (``URLError``, resets, short reads, HTTP 5xx/408/429) are retried with backoff of
  1, 2, 4, ... seconds (capped at 30) like :func:`nida.bench.fetch.download`. A server that
  answers 200 to a ranged request is refused (it would send the whole archive).
* The central directory is trusted for sizes; the local header is cross-checked (signature,
  name, method). The CRC-32 and the size of every extracted member are verified, and the
  decompressed size is bounded by the declared size (no zip bomb).
"""

from __future__ import annotations

import logging
import os
import re
import struct
import tempfile
import time
import urllib.error
import urllib.request
import zlib
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import IO

from nida.bench import fetch

__all__ = ["RemoteZip", "RemoteZipError", "ZipMember", "parse_central_directory"]

log = logging.getLogger(__name__)

_USER_AGENT = "nida-data/0"
_CONTENT_RANGE_RE = re.compile(r"bytes (\d+)-(\d+)/(\d+|\*)")
_MAX_BACKOFF_S = 30.0
_RETRYABLE_4XX = frozenset({408, 429})
# The transient-network set of nida.bench.fetch (same retry classification).
_NETWORK_ERRORS: tuple[type[BaseException], ...] = fetch._NETWORK_ERRORS
_COPY_STEP_BYTES = 1 << 20

_EOCD_SIG = b"PK\x05\x06"
_EOCD_FMT = "<4s4H2LH"  # sig, disk, cd_disk, entries_on_disk, entries, cd_size, cd_offset, comment
_EOCD_SIZE = struct.calcsize(_EOCD_FMT)
_LOCATOR_SIG = b"PK\x06\x07"
_LOCATOR_FMT = "<4sLQL"  # sig, disk_of_eocd64, eocd64_offset, total_disks
_LOCATOR_SIZE = struct.calcsize(_LOCATOR_FMT)
_EOCD64_SIG = b"PK\x06\x06"
_EOCD64_FMT = "<4sQ2H2L4Q"  # sig, size, made, needed, disk, cd_disk, n_disk, n, cd_size, cd_offset
_EOCD64_SIZE = struct.calcsize(_EOCD64_FMT)
_CD_SIG = b"PK\x01\x02"
_CD_FMT = "<4s4B4HL2L5H2L"
_CD_SIZE = struct.calcsize(_CD_FMT)
_LOCAL_SIG = b"PK\x03\x04"
_LOCAL_FMT = "<4s2B4HL2L2H"
_LOCAL_SIZE = struct.calcsize(_LOCAL_FMT)

_U16 = 0xFFFF
_U32 = 0xFFFFFFFF
_ZIP64_EXTRA_ID = 0x0001
_FLAG_ENCRYPTED = 0x0001
_FLAG_STRONG_ENCRYPTION = 0x0040
_FLAG_UTF8 = 0x0800
METHOD_STORED = 0
METHOD_DEFLATED = 8


class RemoteZipError(fetch.FetchError):
    """Reading the remote archive failed (network, unsupported or corrupt zip)."""


class _Transient(Exception):
    """Internal: this request failed but an identical retry may succeed."""


@dataclass(frozen=True)
class ZipMember:
    """One central-directory entry.

    Attributes:
        name: Entry name (``/``-separated; a trailing ``/`` marks a directory).
        raw_name: Name bytes as stored (compared with the local header).
        flags: General purpose bit flags.
        method: Compression method (0 stored, 8 deflate are extractable).
        crc32: CRC-32 of the uncompressed data.
        compressed_size: Bytes of (compressed) data.
        file_size: Bytes after decompression.
        header_offset: Offset of the local file header in the archive.
        end_offset: Offset of the next entry (or of the central directory): everything in
            ``[header_offset, end_offset)`` belongs to this member.
    """

    name: str
    raw_name: bytes
    flags: int
    method: int
    crc32: int
    compressed_size: int
    file_size: int
    header_offset: int
    end_offset: int

    @property
    def is_dir(self) -> bool:
        return self.name.endswith("/")


# --------------------------------------------------------------------------------------
# Pure parsing
# --------------------------------------------------------------------------------------


@dataclass(frozen=True)
class _Directory:
    """Where the central directory is and how many entries it holds."""

    offset: int
    size: int
    n_entries: int


def _find_eocd(tail: bytes) -> int:
    """Index of the end-of-central-directory record in ``tail`` (last valid signature).

    A candidate counts only if its comment length reaches exactly the end of ``tail``, so a
    signature inside the archive comment is skipped.
    """
    end = len(tail)
    while True:
        pos = tail.rfind(_EOCD_SIG, 0, end)
        if pos < 0:
            raise RemoteZipError("end of central directory record not found (not a zip archive?)")
        if pos + _EOCD_SIZE <= len(tail):
            (comment_len,) = struct.unpack_from("<H", tail, pos + _EOCD_SIZE - 2)
            if pos + _EOCD_SIZE + comment_len == len(tail):
                return pos
        end = pos + len(_EOCD_SIG) - 1  # exclude the occurrence at ``pos`` from the next search


def _classic_directory(eocd: bytes) -> tuple[_Directory, bool]:
    """Parse a classic EOCD; second value: True when any field is a ZIP64 sentinel."""
    _, disk, cd_disk, _, n_entries, cd_size, cd_offset, _ = struct.unpack(_EOCD_FMT, eocd)
    sentinel = n_entries == _U16 or cd_size == _U32 or cd_offset == _U32
    if (disk, cd_disk) != (0, 0) and not sentinel:
        raise RemoteZipError(f"multi-disk zip archives are not supported (disk {disk}/{cd_disk})")
    return _Directory(cd_offset, cd_size, n_entries), sentinel


def _zip64_directory(record: bytes) -> _Directory:
    sig, _, _, _, disk, cd_disk, _, n_entries, cd_size, cd_offset = struct.unpack(
        _EOCD64_FMT, record
    )
    if sig != _EOCD64_SIG:
        raise RemoteZipError("ZIP64 end of central directory record has a bad signature")
    if (disk, cd_disk) != (0, 0):
        raise RemoteZipError(f"multi-disk zip archives are not supported (disk {disk}/{cd_disk})")
    return _Directory(cd_offset, cd_size, n_entries)


def _apply_zip64_extra(
    extra: bytes, file_size: int, compressed_size: int, header_offset: int, name: str
) -> tuple[int, int, int]:
    """Replace sentinel (``0xFFFFFFFF``) sizes/offset by the values of the ZIP64 extra field.

    The ZIP64 field holds, in this order, only those of uncompressed size, compressed size
    and header offset whose 32-bit value is the sentinel.
    """
    needed = [file_size == _U32, compressed_size == _U32, header_offset == _U32]
    pos = 0
    while pos + 4 <= len(extra):
        field_id, field_size = struct.unpack_from("<HH", extra, pos)
        body = extra[pos + 4 : pos + 4 + field_size]
        pos += 4 + field_size
        if field_id != _ZIP64_EXTRA_ID:
            continue
        count = sum(needed)
        if len(body) < 8 * count:
            raise RemoteZipError(f"{name!r}: truncated ZIP64 extra field")
        values = iter(struct.unpack_from(f"<{count}Q", body, 0))
        return (
            next(values) if needed[0] else file_size,
            next(values) if needed[1] else compressed_size,
            next(values) if needed[2] else header_offset,
        )
    if any(needed):
        raise RemoteZipError(f"{name!r}: ZIP64 sentinel without a ZIP64 extra field")
    return file_size, compressed_size, header_offset


def parse_central_directory(data: bytes, n_entries: int, directory_offset: int) -> list[ZipMember]:
    """Parse ``n_entries`` central-directory entries from ``data``.

    Args:
        data: The central directory bytes (exactly ``cd_size`` of them).
        n_entries: Entry count from the end record.
        directory_offset: Archive offset of ``data``; the last member ends there.

    Returns:
        Members in directory order, ``end_offset`` filled in.

    Raises:
        RemoteZipError: on a bad signature, truncation, trailing bytes or a count mismatch.
    """
    pos = 0
    rows: list[tuple[str, bytes, int, int, int, int, int, int]] = []
    for index in range(n_entries):
        if pos + _CD_SIZE > len(data):
            raise RemoteZipError(f"central directory truncated at entry {index}/{n_entries}")
        fields = struct.unpack_from(_CD_FMT, data, pos)
        if fields[0] != _CD_SIG:
            raise RemoteZipError(f"central directory entry {index} has a bad signature")
        flags, method = fields[5], fields[6]
        crc, comp, size = fields[9], fields[10], fields[11]
        name_len, extra_len, comment_len = fields[12], fields[13], fields[14]
        offset = fields[18]
        start = pos + _CD_SIZE
        end = start + name_len + extra_len + comment_len
        if end > len(data):
            raise RemoteZipError(f"central directory truncated inside entry {index}")
        raw_name = bytes(data[start : start + name_len])
        extra = bytes(data[start + name_len : start + name_len + extra_len])
        name = raw_name.decode("utf-8" if flags & _FLAG_UTF8 else "cp437", errors="replace")
        size, comp, offset = _apply_zip64_extra(extra, size, comp, offset, name)
        rows.append((name, raw_name, flags, method, crc, comp, size, offset))
        pos = end
    if pos != len(data):
        raise RemoteZipError(f"central directory has {len(data) - pos} unparsed trailing bytes")
    boundaries = sorted({row[7] for row in rows} | {directory_offset})
    members: list[ZipMember] = []
    for name, raw_name, flags, method, crc, comp, size, offset in rows:
        later = [b for b in boundaries if b > offset]
        members.append(
            ZipMember(
                name=name,
                raw_name=raw_name,
                flags=flags,
                method=method,
                crc32=crc,
                compressed_size=comp,
                file_size=size,
                header_offset=offset,
                end_offset=later[0] if later else directory_offset,
            )
        )
    return members


# --------------------------------------------------------------------------------------
# RemoteZip
# --------------------------------------------------------------------------------------


class RemoteZip:
    """A zip archive behind an HTTP(S) URL that supports ``Range`` requests.

    Args:
        url: ``http://`` or ``https://`` URL of the archive.
        opener: Network access function (injectable for tests).
        timeout_s: Socket timeout per request, seconds.
        max_attempts: Attempts per request (first try included), at least 1.
        sleep: Sleep function used for backoff.
        window_bytes: Read-ahead window; every fetch asks for at least this many bytes.
        tail_bytes: Bytes read from the end of the archive to find the end record (must be
            able to hold a record plus the longest comment, 65 557 bytes).
    """

    def __init__(
        self,
        url: str,
        *,
        opener: fetch.Opener = fetch.default_opener,
        timeout_s: float = 60.0,
        max_attempts: int = 5,
        sleep: Callable[[float], None] = time.sleep,
        window_bytes: int = 4 << 20,
        tail_bytes: int = 256 << 10,
    ) -> None:
        if not url.lower().startswith(("http://", "https://")):
            raise ValueError(f"only http(s) URLs are supported, got {url!r}")
        if timeout_s <= 0:
            raise ValueError(f"timeout_s must be > 0, got {timeout_s}")
        if max_attempts < 1:
            raise ValueError(f"max_attempts must be >= 1, got {max_attempts}")
        if window_bytes < 1:
            raise ValueError(f"window_bytes must be >= 1, got {window_bytes}")
        if tail_bytes < _EOCD_SIZE + _LOCATOR_SIZE + 0xFFFF:
            raise ValueError(f"tail_bytes must hold an end record and a comment, got {tail_bytes}")
        self.url = url
        self._opener = opener
        self._timeout_s = timeout_s
        self._max_attempts = max_attempts
        self._sleep = sleep
        self._window_bytes = window_bytes
        self._tail_bytes = tail_bytes
        self._size: int | None = None
        self._members: list[ZipMember] | None = None
        self._window_start = 0
        self._window = b""
        self.num_requests = 0

    # ------------------------------------------------------------------ HTTP

    def _attempt(self, start: int, stop: int) -> tuple[bytes, int | None]:
        request = urllib.request.Request(
            self.url, headers={"User-Agent": _USER_AGENT, "Range": f"bytes={start}-{stop - 1}"}
        )
        self.num_requests += 1
        try:
            response = self._opener(request, self._timeout_s)
        except urllib.error.HTTPError as exc:
            code = exc.code
            exc.close()
            if 400 <= code < 500 and code not in _RETRYABLE_4XX:
                raise RemoteZipError(
                    f"HTTP {code} for {self.url} (bytes {start}-{stop - 1})"
                ) from exc
            raise _Transient(f"HTTP {code}") from exc
        except _NETWORK_ERRORS as exc:
            raise _Transient(f"connecting failed: {exc!r}") from exc
        with response:
            if response.status == 200:
                raise RemoteZipError(
                    f"{self.url} ignored the Range header (HTTP 200); refusing to download "
                    "the whole archive"
                )
            if response.status != 206:
                raise RemoteZipError(f"unexpected HTTP status {response.status} from {self.url}")
            header = response.headers.get("Content-Range")
            match = _CONTENT_RANGE_RE.fullmatch(header.strip()) if header else None
            if match is None:
                raise RemoteZipError(
                    f"missing or malformed Content-Range from {self.url}: {header!r}"
                )
            first, last = int(match.group(1)), int(match.group(2))
            total = None if match.group(3) == "*" else int(match.group(3))
            if first != start or last > stop - 1 or last < first:
                raise RemoteZipError(
                    f"{self.url}: asked for bytes {start}-{stop - 1}, server sent {first}-{last}"
                )
            expected = last - first + 1
            chunks: list[bytes] = []
            received = 0
            try:
                while received < expected:
                    chunk = response.read(min(_COPY_STEP_BYTES, expected - received))
                    if not chunk:
                        break
                    chunks.append(chunk)
                    received += len(chunk)
            except _NETWORK_ERRORS as exc:
                raise _Transient(f"reading failed after {received} bytes: {exc!r}") from exc
        if received != expected:
            raise _Transient(f"short body: {received} bytes, expected {expected}")
        return b"".join(chunks), total

    def _request(self, start: int, stop: int) -> tuple[bytes, int | None]:
        """Bytes ``[start, stop)`` with retries; also the archive size if the server says."""
        last: _Transient | None = None
        for attempt in range(1, self._max_attempts + 1):
            try:
                return self._attempt(start, stop)
            except _Transient as exc:
                last = exc
                log.warning(
                    "range %d-%d of %s: attempt %d/%d failed: %s",
                    start,
                    stop - 1,
                    self.url,
                    attempt,
                    self._max_attempts,
                    exc,
                )
                if attempt < self._max_attempts:
                    self._sleep(min(2.0 ** (attempt - 1), _MAX_BACKOFF_S))
        raise RemoteZipError(
            f"range request {start}-{stop - 1} of {self.url} failed after "
            f"{self._max_attempts} attempts: {last}"
        ) from last

    @property
    def size_bytes(self) -> int:
        """Archive size, learned from the ``Content-Range`` of a one-byte probe."""
        if self._size is None:
            _, total = self._request(0, 1)
            if total is None or total < _EOCD_SIZE:
                raise RemoteZipError(f"cannot determine the size of {self.url} (got {total!r})")
            self._size = total
        return self._size

    def _read(self, start: int, length: int) -> bytes:
        """``length`` bytes at ``start``, served from the read-ahead window when possible."""
        size = self.size_bytes
        if length < 0 or start < 0 or start + length > size:
            raise RemoteZipError(
                f"{self.url}: read of {length} bytes at {start} is outside the archive ({size})"
            )
        if length == 0:
            return b""
        offset = start - self._window_start
        if offset >= 0 and offset + length <= len(self._window):
            return self._window[offset : offset + length]
        stop = min(size, start + max(length, self._window_bytes))
        body, _ = self._request(start, stop)
        if len(body) != stop - start:
            raise RemoteZipError(
                f"{self.url}: server returned {len(body)} bytes for a {stop - start}-byte range"
            )
        self._window_start, self._window = start, body
        return body[:length]

    # ------------------------------------------------------------------ listing

    def _locate_directory(self) -> tuple[_Directory, int]:
        """Central directory location and the archive offset where its end record begins."""
        size = self.size_bytes
        tail_len = min(size, self._tail_bytes)
        tail = self._read(size - tail_len, tail_len)
        eocd_at = _find_eocd(tail)
        eocd_pos = size - tail_len + eocd_at
        directory, sentinel = _classic_directory(tail[eocd_at : eocd_at + _EOCD_SIZE])
        locator_at = eocd_at - _LOCATOR_SIZE
        has_locator = locator_at >= 0 and tail[locator_at : locator_at + 4] == _LOCATOR_SIG
        if not has_locator:
            if sentinel:
                raise RemoteZipError("ZIP64 sentinel values but no ZIP64 end locator")
            return directory, eocd_pos
        _, locator_disk, eocd64_offset, total_disks = struct.unpack_from(
            _LOCATOR_FMT, tail, locator_at
        )
        if locator_disk != 0 or total_disks > 1:
            raise RemoteZipError("multi-disk zip archives are not supported")
        if eocd64_offset + _EOCD64_SIZE > size:
            raise RemoteZipError("ZIP64 end record offset lies outside the archive")
        directory = _zip64_directory(self._read(eocd64_offset, _EOCD64_SIZE))
        return directory, eocd64_offset

    def members(self) -> list[ZipMember]:
        """All central-directory entries (directories included), in directory order.

        The result is cached for the lifetime of the object.

        Raises:
            RemoteZipError: on network failure or an unsupported or corrupt archive.
        """
        if self._members is None:
            directory, end_limit = self._locate_directory()
            if directory.offset + directory.size > end_limit:
                raise RemoteZipError(
                    f"central directory ({directory.offset}+{directory.size}) overlaps the end "
                    f"record at {end_limit}; archives with prepended data are not supported"
                )
            data = self._read(directory.offset, directory.size)
            self._members = parse_central_directory(data, directory.n_entries, directory.offset)
            log.info("%s: %d entries, %d bytes", self.url, len(self._members), self.size_bytes)
        return list(self._members)

    # ------------------------------------------------------------------ extraction

    def extract(self, member: ZipMember, dest: Path, *, skip_existing: bool = True) -> Path:
        """Write one member to ``dest`` (atomic) and verify its size and CRC-32.

        Args:
            member: An entry of :meth:`members`.
            dest: Destination file; parent folders are created.
            skip_existing: Return ``dest`` without any network access when it already has the
                member's size (files only ever appear complete, see the module docstring).

        Returns:
            ``dest``.

        Raises:
            RemoteZipError: for a directory, an encrypted or unsupported member, a local
                header that contradicts the central directory, corrupt data, or network
                failure.
        """
        if member.is_dir:
            raise RemoteZipError(f"{member.name!r} is a directory")
        if member.flags & (_FLAG_ENCRYPTED | _FLAG_STRONG_ENCRYPTION):
            raise RemoteZipError(f"{member.name!r} is encrypted")
        if member.method not in (METHOD_STORED, METHOD_DEFLATED):
            raise RemoteZipError(f"{member.name!r}: unsupported compression method {member.method}")
        if skip_existing and dest.is_file() and dest.stat().st_size == member.file_size:
            return dest
        data_start = self._data_start(member)
        dest.parent.mkdir(parents=True, exist_ok=True)
        fd, tmp_name = tempfile.mkstemp(prefix=f".{dest.name}.", suffix=".tmp", dir=dest.parent)
        tmp_path = Path(tmp_name)
        try:
            with os.fdopen(fd, "wb") as handle:
                self._copy_data(member, data_start, handle)
                handle.flush()
                os.fsync(handle.fileno())
            tmp_path.replace(dest)
        except BaseException:
            tmp_path.unlink(missing_ok=True)
            raise
        return dest

    def _data_start(self, member: ZipMember) -> int:
        """Archive offset of the member's data, after checking its local header."""
        header = self._read(member.header_offset, _LOCAL_SIZE)
        fields = struct.unpack(_LOCAL_FMT, header)
        if fields[0] != _LOCAL_SIG:
            raise RemoteZipError(f"{member.name!r}: bad local header signature")
        method, name_len, extra_len = int(fields[4]), int(fields[10]), int(fields[11])
        if method != member.method:
            raise RemoteZipError(
                f"{member.name!r}: local header method {method} != central directory "
                f"{member.method}"
            )
        local_name = self._read(member.header_offset + _LOCAL_SIZE, name_len)
        if local_name != member.raw_name:
            raise RemoteZipError(f"{member.name!r}: local header names {local_name!r}")
        data_start = member.header_offset + _LOCAL_SIZE + name_len + extra_len
        if data_start + member.compressed_size > member.end_offset:
            raise RemoteZipError(
                f"{member.name!r}: data ({data_start}+{member.compressed_size}) runs into the "
                f"next entry at {member.end_offset}"
            )
        return data_start

    def _copy_data(self, member: ZipMember, data_start: int, handle: IO[bytes]) -> None:
        write = handle.write
        inflater = zlib.decompressobj(-zlib.MAX_WBITS) if member.method == METHOD_DEFLATED else None
        crc = 0
        written = 0
        position = data_start
        remaining = member.compressed_size
        while remaining > 0:
            chunk = self._read(position, min(remaining, self._window_bytes))
            position += len(chunk)
            remaining -= len(chunk)
            pending = chunk
            while pending:
                if inflater is None:
                    out, pending = pending, b""
                else:
                    try:
                        out = inflater.decompress(pending, _COPY_STEP_BYTES)
                    except zlib.error as exc:
                        raise RemoteZipError(
                            f"{member.name!r}: corrupt deflate data: {exc}"
                        ) from exc
                    pending = inflater.unconsumed_tail
                written += len(out)
                if written > member.file_size:
                    raise RemoteZipError(
                        f"{member.name!r}: more than the declared {member.file_size} bytes"
                    )
                crc = zlib.crc32(out, crc)
                write(out)
        if inflater is not None:
            try:
                tail = inflater.flush()
            except zlib.error as exc:
                raise RemoteZipError(f"{member.name!r}: corrupt deflate data: {exc}") from exc
            written += len(tail)
            crc = zlib.crc32(tail, crc)
            write(tail)
            if not inflater.eof:
                raise RemoteZipError(f"{member.name!r}: deflate stream is truncated")
        if written != member.file_size:
            raise RemoteZipError(
                f"{member.name!r}: {written} bytes extracted, expected {member.file_size}"
            )
        if crc != member.crc32:
            raise RemoteZipError(
                f"{member.name!r}: CRC-32 {crc:08x} does not match the archive's {member.crc32:08x}"
            )