backend/archive_utils.py
Rama master · réplica de lectura del repositorio original.
"""Shared archive extraction utilities with ZIP bomb and size protection.
All subtitle providers that handle archive downloads should use these helpers
instead of inline zipfile/rarfile calls, to ensure consistent security checks
across the entire download pipeline.
"""
import contextlib
import hashlib
import io
import logging
import os
import zipfile
from werkzeug.utils import secure_filename
logger = logging.getLogger(__name__)
_MAX_ARCHIVE_BYTES = 20 * 1024 * 1024 # 20 MB — reject before extraction
_MAX_EXTRACTED_BYTES = 50 * 1024 * 1024 # 50 MB — total extracted
_MAX_COMPRESSION_RATIO = 100 # 100:1 — ZIP bomb ratio
_SUBTITLE_EXTENSIONS = frozenset({".ass", ".srt", ".ssa", ".vtt"})
def _safe_basename(path: str) -> str:
"""Return a filename-safe basename (P2).
``os.path.basename`` is platform-dependent: on Linux it does not strip
backslashes, so a ZIP entry named ``..\\..\\cmd.exe.srt`` would slip
through on non-Windows hosts. ``secure_filename`` normalises to ASCII,
strips traversal components, and blocks control chars / nul / slashes
on every platform.
"""
raw = os.path.basename(path.replace("\\", "/"))
safe = secure_filename(raw)
stem, ext = os.path.splitext(raw)
if ext and not safe.lower().endswith(ext.lower()) and secure_filename(ext.lstrip(".")):
# Every letter of the name was non-ASCII ("简体.srt" -> "srt"): keep the
# extension, or the subtitle check drops the member, and make the name
# unique, or two members collapse into one.
digest = hashlib.sha1(
raw.encode("utf-8", "surrogateescape"), usedforsecurity=False
).hexdigest()[:8]
safe_stem = secure_filename(stem) or "subtitle"
return f"{safe_stem}-{digest}{ext.lower()}"
return safe
def _plain_basename(path: str) -> str:
"""Basename that keeps non-Latin letters, still safe as a label.
``secure_filename`` reduces "简体.srt" to "srt" and "Ταινία.srt" to "srt":
the extension check then drops the file, and two members can collapse to
one name. Callers that only match or display names (never write them as a
path) use this instead: directories, control characters and leading dots
are removed, letters of every script stay.
"""
raw = os.path.basename(path.replace("\\", "/"))
cleaned = "".join(ch for ch in raw if ch.isprintable()).strip().lstrip(".")
return cleaned
def extract_subtitles_from_zip(
data: bytes,
subtitle_exts: frozenset[str] = _SUBTITLE_EXTENSIONS,
raw_names: bool = False,
) -> list[tuple[str, bytes]]:
"""In-memory ZIP extraction with ZIP bomb and size protection.
Args:
data: Raw ZIP archive bytes.
subtitle_exts: Allowed subtitle file extensions (with leading dot).
raw_names: Return names via ``_plain_basename`` (non-Latin letters kept)
instead of ``secure_filename``. Only for callers that never use the
name as a filesystem path.
Returns:
List of (basename, content) tuples for matching subtitle files.
Raises:
ValueError: If archive exceeds size limits or compression ratio.
"""
if len(data) > _MAX_ARCHIVE_BYTES:
raise ValueError(
f"Archive too large: {len(data) // (1024 * 1024)} MB > "
f"{_MAX_ARCHIVE_BYTES // (1024 * 1024)} MB limit"
)
try:
with zipfile.ZipFile(io.BytesIO(data)) as zf:
total_compressed = sum(info.compress_size for info in zf.infolist())
total_uncompressed = sum(info.file_size for info in zf.infolist())
if total_uncompressed > _MAX_EXTRACTED_BYTES:
raise ValueError(
f"Archive extraction too large: {total_uncompressed // (1024 * 1024)} MB > "
f"{_MAX_EXTRACTED_BYTES // (1024 * 1024)} MB limit"
)
# ZIP bomb ratio check (skip if nothing compressed — empty/stored zips)
if total_compressed > 0:
ratio = total_uncompressed / total_compressed
if ratio > _MAX_COMPRESSION_RATIO:
raise ValueError(
f"ZIP bomb detected: compression ratio {ratio:.0f}:1 exceeds "
f"{_MAX_COMPRESSION_RATIO}:1 limit"
)
results = []
for info in zf.infolist():
# Skip directory entries
if info.filename.endswith("/"):
continue
# Strip any path components + sanitize for disk write (ZIP slip + P2)
basename = (_plain_basename if raw_names else _safe_basename)(info.filename)
if not basename:
continue
ext = os.path.splitext(basename)[1].lower()
if ext not in subtitle_exts:
continue
results.append((basename, zf.read(info.filename)))
return results
except zipfile.BadZipFile:
logger.warning("archive_utils: bad ZIP archive, skipping")
return []
def safe_read_zip_member(
zf: zipfile.ZipFile,
name: str,
max_bytes: int = _MAX_EXTRACTED_BYTES,
) -> bytes:
"""Read a single named entry from an open ZipFile with bomb protection.
Use when callers know which entries they want (e.g. ``config.json``)
and need to avoid loading a multi-GB decompressed payload into memory.
The caller has already validated the outer archive size; this guards
each individual entry against a single oversized member.
Args:
zf: Open ``zipfile.ZipFile`` instance.
name: Entry name to read.
max_bytes: Cap on the entry's uncompressed size (default 50 MB,
shared with ``extract_subtitles_from_zip``).
Returns:
Decompressed entry bytes.
Raises:
KeyError: If ``name`` is not in the archive.
ValueError: If the uncompressed size or compression ratio
exceeds the configured limits.
"""
info = zf.getinfo(name) # raises KeyError if missing
# The central-directory header fields are attacker-controlled: a crafted
# archive can declare file_size=0 while storing a large DEFLATE payload,
# making both header checks pass and then expanding to GBs inside zf.read().
# Treat the headers only as a cheap fast-reject, and enforce the real cap by
# streaming the decompressed bytes and stopping one byte past the limit.
if info.file_size > max_bytes:
raise ValueError(
f"ZIP entry '{name}' too large: {info.file_size // (1024 * 1024)} MB > "
f"{max_bytes // (1024 * 1024)} MB limit"
)
if info.compress_size > 0:
ratio = info.file_size / info.compress_size
if ratio > _MAX_COMPRESSION_RATIO:
raise ValueError(
f"ZIP bomb detected on '{name}': ratio {ratio:.0f}:1 exceeds "
f"{_MAX_COMPRESSION_RATIO}:1 limit"
)
with zf.open(name) as member:
data = member.read(max_bytes + 1)
if len(data) > max_bytes:
raise ValueError(
f"ZIP entry '{name}' exceeds {max_bytes // (1024 * 1024)} MB cap during decompression"
)
return data
def safe_extract_zip_member_to(
zf: zipfile.ZipFile,
name: str,
dest_path: str,
max_bytes: int = _MAX_EXTRACTED_BYTES,
) -> int:
"""Stream one entry to ``dest_path`` with the same guards as the reader.
The in-memory reader is fine for a config file, but a database inside a
full backup can be gigabytes: materialising it (on top of the archive
itself) is a memory-exhaustion risk. Returns the bytes written; the partial
file is removed when a limit is hit.
"""
info = zf.getinfo(name) # raises KeyError if missing
if info.file_size > max_bytes:
raise ValueError(
f"ZIP entry '{name}' too large: {info.file_size // (1024 * 1024)} MB > "
f"{max_bytes // (1024 * 1024)} MB limit"
)
if info.compress_size > 0:
ratio = info.file_size / info.compress_size
if ratio > _MAX_COMPRESSION_RATIO:
raise ValueError(
f"ZIP bomb detected on '{name}': ratio {ratio:.0f}:1 exceeds "
f"{_MAX_COMPRESSION_RATIO}:1 limit"
)
written = 0
try:
with zf.open(name) as member, open(dest_path, "wb") as out:
while chunk := member.read(1024 * 1024):
written += len(chunk)
if written > max_bytes:
raise ValueError(
f"ZIP entry '{name}' exceeds {max_bytes // (1024 * 1024)} MB cap "
"during decompression"
)
out.write(chunk)
except BaseException:
with contextlib.suppress(OSError):
os.remove(dest_path)
raise
return written
def extract_subtitles_from_rar(
data: bytes,
subtitle_exts: frozenset[str] = _SUBTITLE_EXTENSIONS,
raw_names: bool = False,
) -> list[tuple[str, bytes]]:
"""In-memory RAR extraction with size protection.
Args:
data: Raw RAR archive bytes.
subtitle_exts: Allowed subtitle file extensions (with leading dot).
Returns:
List of (basename, content) tuples for matching subtitle files.
Raises:
ImportError: If the rarfile package is not installed.
ValueError: If archive exceeds the size limit.
"""
if len(data) > _MAX_ARCHIVE_BYTES:
raise ValueError(
f"Archive too large: {len(data) // (1024 * 1024)} MB > "
f"{_MAX_ARCHIVE_BYTES // (1024 * 1024)} MB limit"
)
import rarfile # optional dep — raises ImportError if absent
results = []
try:
with rarfile.RarFile(io.BytesIO(data)) as rf:
for info in rf.infolist():
basename = (_plain_basename if raw_names else _safe_basename)(info.filename)
if not basename:
continue
ext = os.path.splitext(basename)[1].lower()
if ext not in subtitle_exts:
continue
results.append((basename, rf.read(info.filename)))
except Exception as e:
logger.warning("archive_utils: RAR extraction failed: %s", e)
return results