Files
govoplan-files/src/govoplan_files/backend/storage/archives.py
T

766 lines
27 KiB
Python

from __future__ import annotations
import mimetypes
import re
import stat
import tarfile
import zipfile
from dataclasses import dataclass
from io import BytesIO
from os import PathLike
from pathlib import Path, PurePosixPath
from typing import Any, BinaryIO, Callable, Iterable, Iterator, Literal
import pyzipper
from sqlalchemy.orm import Session
from govoplan_files.backend.db.models import FileAsset
from govoplan_files.backend.storage.backends import (
StorageBackendError,
get_storage_backend,
)
from govoplan_files.backend.storage.common import (
FileConflictResolution,
FileStorageError,
UploadedStoredFile,
)
from govoplan_files.backend.storage.files import (
create_file_asset,
current_versions_and_blobs,
)
from govoplan_files.backend.storage.paths import (
filename_from_path,
normalize_folder,
normalize_logical_path,
)
from govoplan_files.backend.storage.native_zip import (
native_zip_library,
read_native_zip_members,
)
_ARCHIVE_READ_CHUNK_SIZE = 1024 * 1024
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")
ARCHIVE_UPLOAD_MAX_ENTRIES = 10_000
ARCHIVE_MAX_PATH_BYTES = 4096
ARCHIVE_MAX_PATH_DEPTH = 128
# Kept for callers and documentation using the former ZIP-specific name.
ZIP_UPLOAD_MAX_FILES = ARCHIVE_UPLOAD_MAX_ENTRIES
SUPPORTED_ARCHIVE_SUFFIXES = (
".tar.bz2",
".tar.gz",
".tar.xz",
".tbz2",
".tgz",
".txz",
".tar",
".zip",
)
ArchiveProgress = Callable[[str, int, int, int, int], None]
class ArchivePasswordError(FileStorageError):
pass
@dataclass(frozen=True, slots=True)
class ArchiveEntry:
path: str
kind: Literal["file", "directory"]
size_bytes: int
compressed_size_bytes: int | None = None
encrypted: bool = False
@dataclass(frozen=True, slots=True)
class ArchiveInspection:
archive_format: str
entries: tuple[ArchiveEntry, ...]
file_count: int
directory_count: int
expanded_size_bytes: int
compressed_size_bytes: int
requires_password: bool
password_verified: bool
def create_zip_file(
session: Session, assets: Iterable[FileAsset], output_path: str | Path
) -> None:
backend = get_storage_backend()
asset_list = list(assets)
version_blobs = current_versions_and_blobs(session, asset_list)
with zipfile.ZipFile(
output_path, mode="w", compression=zipfile.ZIP_DEFLATED
) as archive:
for asset in asset_list:
_version, blob = version_blobs[asset.id]
info = zipfile.ZipInfo(asset.display_path)
info.compress_type = zipfile.ZIP_DEFLATED
with archive.open(info, "w") as member:
try:
for chunk in backend.iter_bytes(blob.storage_key):
if chunk:
member.write(chunk)
except StorageBackendError as exc:
raise FileStorageError(str(exc)) from exc
def archive_format_for_filename(filename: str) -> str:
lowered = filename.strip().casefold()
if lowered.endswith(".zip"):
return "zip"
if lowered.endswith((".tar.gz", ".tgz")):
return "tar.gz"
if lowered.endswith((".tar.bz2", ".tbz2")):
return "tar.bz2"
if lowered.endswith((".tar.xz", ".txz")):
return "tar.xz"
if lowered.endswith(".tar"):
return "tar"
raise FileStorageError(
"Unsupported archive format. Use ZIP, TAR, TAR.GZ, TAR.BZ2, or TAR.XZ."
)
def is_supported_archive_filename(filename: str) -> bool:
try:
archive_format_for_filename(filename)
except FileStorageError:
return False
return True
def inspect_archive(
archive_data: bytes | str | PathLike[str],
*,
filename: str,
password: str | None = None,
max_entries: int = ARCHIVE_UPLOAD_MAX_ENTRIES,
max_expanded_bytes: int = 2 * 1024 * 1024 * 1024,
max_expansion_ratio: int = 100,
) -> ArchiveInspection:
from govoplan_files.backend.storage.archive_workers import inspect_archive_isolated
return inspect_archive_isolated(
archive_data, filename=filename, password=password, max_entries=max_entries,
max_expanded_bytes=max_expanded_bytes, max_expansion_ratio=max_expansion_ratio,
)
def _inspect_archive_content(
archive_data: bytes | str | PathLike[str],
*,
filename: str,
password: str | None = None,
max_entries: int = ARCHIVE_UPLOAD_MAX_ENTRIES,
max_expanded_bytes: int = 2 * 1024 * 1024 * 1024,
max_expansion_ratio: int = 100,
) -> ArchiveInspection:
archive_format = archive_format_for_filename(filename)
compressed_size = _archive_size(archive_data)
if archive_format == "zip":
entries, requires_password, password_verified = _inspect_zip(
archive_data,
password=password,
)
else:
entries = _inspect_tar(
archive_data,
compressed_size=compressed_size,
max_entries=max_entries,
max_expanded_bytes=max_expanded_bytes,
max_expansion_ratio=max_expansion_ratio,
)
requires_password = False
password_verified = True
_validate_archive_limits(
entries,
compressed_size=compressed_size,
max_entries=max_entries,
max_expanded_bytes=max_expanded_bytes,
max_expansion_ratio=max_expansion_ratio,
)
complete_entries = _with_derived_directories(entries, max_entries=max_entries)
return ArchiveInspection(
archive_format=archive_format,
entries=tuple(complete_entries),
file_count=sum(entry.kind == "file" for entry in complete_entries),
directory_count=sum(entry.kind == "directory" for entry in complete_entries),
expanded_size_bytes=sum(
entry.size_bytes for entry in entries if entry.kind == "file"
),
compressed_size_bytes=compressed_size,
requires_password=requires_password,
password_verified=password_verified,
)
def extract_archive_upload(
session: Session,
*,
tenant_id: str,
owner_type: str,
owner_id: str,
user_id: str,
archive_data: bytes | str | PathLike[str],
filename: str,
folder: str | None,
campaign_id: str | None,
selected_paths: Iterable[str] | None = None,
password: str | None = None,
conflict_strategy: str = "reject",
conflict_resolutions: Iterable[FileConflictResolution] | None = None,
metadata: dict[str, Any] | None = None,
is_admin: bool = False,
encryption_vault_id: str | None = None,
max_entries: int = ARCHIVE_UPLOAD_MAX_ENTRIES,
max_file_bytes: int = 50 * 1024 * 1024,
max_expanded_bytes: int = 2 * 1024 * 1024 * 1024,
max_expansion_ratio: int = 100,
progress: ArchiveProgress | None = None,
) -> list[UploadedStoredFile]:
from govoplan_files.backend.storage.archive_workers import extract_archive_isolated
return extract_archive_isolated(
session, archive_data=archive_data, filename=filename, password=password,
selected_paths=selected_paths, max_entries=max_entries, max_file_bytes=max_file_bytes,
max_expanded_bytes=max_expanded_bytes, max_expansion_ratio=max_expansion_ratio,
progress=progress, store_options={
"tenant_id": tenant_id, "owner_type": owner_type, "owner_id": owner_id,
"user_id": user_id, "folder": folder, "campaign_id": campaign_id,
"conflict_strategy": conflict_strategy, "conflict_resolutions": conflict_resolutions,
"metadata": metadata, "is_admin": is_admin, "encryption_vault_id": encryption_vault_id,
},
)
def extract_zip_upload(
session: Session,
*,
tenant_id: str,
owner_type: str,
owner_id: str,
user_id: str,
zip_data: bytes | str | PathLike[str],
folder: str | None,
campaign_id: str | None,
conflict_strategy: str = "reject",
conflict_resolutions: Iterable[FileConflictResolution] | None = None,
metadata: dict[str, Any] | None = None,
is_admin: bool = False,
encryption_vault_id: str | None = None,
max_files: int = ZIP_UPLOAD_MAX_FILES,
max_file_bytes: int = 50 * 1024 * 1024,
max_total_bytes: int = 250 * 1024 * 1024,
) -> list[UploadedStoredFile]:
"""Backward-compatible wrapper for the original immediate ZIP endpoint."""
return extract_archive_upload(
session,
tenant_id=tenant_id,
owner_type=owner_type,
owner_id=owner_id,
user_id=user_id,
archive_data=zip_data,
filename="archive.zip",
folder=folder,
campaign_id=campaign_id,
conflict_strategy=conflict_strategy,
conflict_resolutions=conflict_resolutions,
metadata=metadata,
is_admin=is_admin,
encryption_vault_id=encryption_vault_id,
max_entries=max_files,
max_file_bytes=max_file_bytes,
max_expanded_bytes=max_total_bytes,
)
def _inspect_zip(
archive_data: bytes | str | PathLike[str],
*,
password: str | None,
) -> tuple[list[ArchiveEntry], bool, bool]:
try:
with pyzipper.AESZipFile(_archive_source(archive_data)) as archive:
infos = archive.infolist()
entries = [_zip_entry(info) for info in infos]
encrypted_info = next(
(
info
for info in infos
if not info.is_dir() and bool(info.flag_bits & 0x1)
),
None,
)
if encrypted_info is None:
return entries, False, True
if not password:
return entries, True, False
try:
with archive.open(
encrypted_info, pwd=password.encode("utf-8")
) as member:
member.read(1)
except (RuntimeError, ValueError, zipfile.BadZipFile) as exc:
raise ArchivePasswordError("Archive password is incorrect") from exc
return entries, True, True
except ArchivePasswordError:
raise
except (OSError, ValueError, zipfile.BadZipFile) as exc:
raise FileStorageError("Invalid ZIP upload") from exc
def _inspect_tar(
archive_data: bytes | str | PathLike[str],
*,
compressed_size: int,
max_entries: int,
max_expanded_bytes: int,
max_expansion_ratio: int,
) -> list[ArchiveEntry]:
try:
with _open_tar(archive_data) as archive:
entries: list[ArchiveEntry] = []
expanded_size = 0
for member in archive:
# Check each header before advancing across its payload. In a
# compressed TAR getmembers() would inflate everything first.
if len(entries) >= max_entries:
raise FileStorageError(
f"Archive contains too many entries (limit {max_entries})"
)
path = _safe_member_path(member.name)
if member.isdir():
entries.append(
ArchiveEntry(path=path, kind="directory", size_bytes=0)
)
continue
if not member.isfile():
raise FileStorageError(
f"Archive member {member.name!r} is not a regular file or directory"
)
if member.size < 0:
raise FileStorageError("Archive member has invalid size metadata")
expanded_size += member.size
if expanded_size > max_expanded_bytes:
raise FileStorageError(
"Archive is too large after extraction "
f"(limit {max_expanded_bytes} bytes)"
)
if expanded_size and (
compressed_size <= 0
or expanded_size > compressed_size * max_expansion_ratio
):
raise FileStorageError(
f"Archive expansion ratio exceeds {max_expansion_ratio}:1"
)
entries.append(
ArchiveEntry(
path=path,
kind="file",
size_bytes=max(0, int(member.size)),
)
)
return entries
except FileStorageError:
raise
except (OSError, tarfile.TarError) as exc:
raise FileStorageError("Invalid TAR upload") from exc
def _zip_entry(info: zipfile.ZipInfo) -> ArchiveEntry:
path = _safe_member_path(info.filename)
unix_mode = (info.external_attr >> 16) & 0xFFFF
file_type = stat.S_IFMT(unix_mode)
if file_type and not (stat.S_ISREG(unix_mode) or stat.S_ISDIR(unix_mode)):
raise FileStorageError(
f"Archive member {info.filename!r} is not a regular file or directory"
)
if info.file_size < 0 or info.compress_size < 0:
raise FileStorageError(
f"Archive member {info.filename!r} has invalid size metadata"
)
return ArchiveEntry(
path=path,
kind="directory" if info.is_dir() else "file",
size_bytes=0 if info.is_dir() else int(info.file_size),
compressed_size_bytes=(None if info.is_dir() else int(info.compress_size)),
encrypted=bool(info.flag_bits & 0x1),
)
def _validate_archive_limits(
entries: list[ArchiveEntry],
*,
compressed_size: int,
max_entries: int,
max_expanded_bytes: int,
max_expansion_ratio: int,
) -> None:
if len(entries) > max_entries:
raise FileStorageError(
f"Archive contains too many entries (limit {max_entries})"
)
seen: dict[str, str] = {}
expanded_size = 0
for entry in entries:
previous_kind = seen.get(entry.path)
if previous_kind is not None:
raise FileStorageError(f"Archive contains duplicate path {entry.path!r}")
seen[entry.path] = entry.kind
if entry.kind == "file":
expanded_size += entry.size_bytes
if expanded_size > max_expanded_bytes:
raise FileStorageError(
"Archive is too large after extraction "
f"(limit {max_expanded_bytes} bytes)"
)
if expanded_size and (
compressed_size <= 0 or expanded_size > compressed_size * max_expansion_ratio
):
raise FileStorageError(
f"Archive expansion ratio exceeds {max_expansion_ratio}:1"
)
def _with_derived_directories(
entries: list[ArchiveEntry],
*,
max_entries: int,
) -> list[ArchiveEntry]:
by_path = {entry.path: entry for entry in entries}
for entry in entries:
parent = PurePosixPath(entry.path).parent
while str(parent) not in {"", "."}:
path = str(parent)
existing = by_path.get(path)
if existing and existing.kind != "directory":
raise FileStorageError(
f"Archive path {path!r} is both a file and a directory"
)
if path not in by_path:
if len(by_path) >= max_entries:
raise FileStorageError(
"Archive contains too many entries including parent "
f"directories (limit {max_entries})"
)
by_path[path] = ArchiveEntry(path=path, kind="directory", size_bytes=0)
parent = parent.parent
return sorted(
by_path.values(),
key=lambda entry: (
tuple(entry.path.casefold().split("/")),
entry.kind != "directory",
),
)
def _selected_file_paths(
entries: tuple[ArchiveEntry, ...],
selected_paths: Iterable[str] | None,
) -> set[str]:
file_paths = {entry.path for entry in entries if entry.kind == "file"}
if selected_paths is None:
return file_paths
known = {entry.path: entry for entry in entries}
normalized = {_safe_member_path(path) for path in selected_paths}
unknown = normalized - known.keys()
if unknown:
raise FileStorageError(
f"Archive selection contains unknown path {sorted(unknown)[0]!r}"
)
selected_files: set[str] = set()
for path in normalized:
entry = known[path]
if entry.kind == "file":
selected_files.add(path)
continue
prefix = f"{path}/"
selected_files.update(
file_path for file_path in file_paths if file_path.startswith(prefix)
)
return selected_files
def _read_selected_zip_members(
archive_data: bytes | str | PathLike[str],
*,
selected_files: set[str],
password: str | None,
max_file_bytes: int,
max_total_bytes: int,
progress: ArchiveProgress | None = None,
total_bytes: int = 0,
) -> Iterator[tuple[str, bytes]]:
total = 0
completed = 0
try:
with pyzipper.AESZipFile(_archive_source(archive_data)) as archive:
infos = archive.infolist()
# Accelerate only classic ZIPCrypto using established public C
# decoding APIs; AES and other formats retain their existing path.
native_compatible = (
bool(password)
and "\x00" not in password
and any(info.flag_bits & 1 for info in infos)
and all(
not getattr(info, "wz_aes_version", None)
and info.compress_type in {zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED}
and (info.filename.isascii() or bool(info.flag_bits & 0x800))
for info in infos
)
)
native = native_zip_library() if native_compatible else None
if native is not None:
yield from read_native_zip_members(
native,
archive_data,
infos_by_path={
_safe_member_path(info.filename): info for info in infos
},
selected_files=selected_files,
normalize_name=_safe_member_path,
password=password,
max_file_bytes=max_file_bytes,
max_total_bytes=max_total_bytes,
progress=progress,
total_bytes=total_bytes,
)
return
for info in infos:
if info.is_dir():
continue
path = _safe_member_path(info.filename)
if path not in selected_files:
continue
pwd = password.encode("utf-8") if password else None
try:
with archive.open(info, pwd=pwd) as source:
data, total = _read_member(
source,
path=path,
max_file_bytes=max_file_bytes,
max_total_bytes=max_total_bytes,
current_total=total,
on_bytes=(
lambda count: progress(
"extracting",
completed,
len(selected_files),
count,
total_bytes,
)
)
if progress
else None,
)
except (RuntimeError, ValueError, zipfile.BadZipFile) as exc:
if info.flag_bits & 0x1:
raise ArchivePasswordError(
"Archive password is incorrect"
) from exc
raise
completed += 1
if progress:
progress(
"extracting", completed, len(selected_files), total, total_bytes
)
yield path, data
del data
except (FileStorageError, ArchivePasswordError):
raise
except (OSError, ValueError, zipfile.BadZipFile) as exc:
raise FileStorageError("ZIP extraction failed") from exc
def _read_selected_tar_members(
archive_data: bytes | str | PathLike[str],
*,
selected_files: set[str],
max_file_bytes: int,
max_total_bytes: int,
progress: ArchiveProgress | None = None,
total_bytes: int = 0,
) -> Iterator[tuple[str, bytes]]:
total = 0
completed = 0
try:
with _open_tar(archive_data) as archive:
# Iterate as headers are read instead of inflating the entire TAR
# with getmembers(), then seeking backwards to inflate it again.
for member in archive:
if not member.isfile():
continue
path = _safe_member_path(member.name)
if path not in selected_files:
continue
source = archive.extractfile(member)
if source is None:
raise FileStorageError(f"Archive member {path!r} could not be read")
with source:
data, total = _read_member(
source,
path=path,
max_file_bytes=max_file_bytes,
max_total_bytes=max_total_bytes,
current_total=total,
on_bytes=(
lambda count: progress(
"extracting",
completed,
len(selected_files),
count,
total_bytes,
)
)
if progress
else None,
)
completed += 1
if progress:
progress(
"extracting", completed, len(selected_files), total, total_bytes
)
yield path, data
del data
except FileStorageError:
raise
except (OSError, tarfile.TarError) as exc:
raise FileStorageError("TAR extraction failed") from exc
def _read_member(
source: BinaryIO,
*,
path: str,
max_file_bytes: int,
max_total_bytes: int,
current_total: int,
on_bytes: Callable[[int], None] | None = None,
) -> tuple[bytes, int]:
parts: list[bytes] = []
actual_size = 0
while True:
read_size = min(
_ARCHIVE_READ_CHUNK_SIZE,
max_file_bytes + 1 - actual_size,
)
chunk = source.read(read_size)
if not chunk:
break
actual_size += len(chunk)
if actual_size > max_file_bytes:
raise FileStorageError(f"Archive member {path!r} exceeds per-file limit")
if current_total + actual_size > max_total_bytes:
raise FileStorageError("Archive is too large after extraction")
parts.append(chunk)
if on_bytes:
on_bytes(current_total + actual_size)
return b"".join(parts), current_total + actual_size
def _store_archive_members(
session: Session,
*,
members: Iterable[tuple[str, bytes]],
tenant_id: str,
owner_type: str,
owner_id: str,
user_id: str,
folder: str | None,
campaign_id: str | None,
conflict_strategy: str,
conflict_resolutions: Iterable[FileConflictResolution] | None,
metadata: dict[str, Any] | None,
is_admin: bool,
encryption_vault_id: str | None,
progress: ArchiveProgress | None = None,
total_files: int = 0,
total_bytes: int = 0,
) -> list[UploadedStoredFile]:
uploaded: list[UploadedStoredFile] = []
base_folder = normalize_folder(folder)
stored_bytes = 0
for inner_path, data in members:
target_path = f"{base_folder}/{inner_path}" if base_folder else inner_path
uploaded.append(
create_file_asset(
session,
tenant_id=tenant_id,
owner_type=owner_type,
owner_id=owner_id,
user_id=user_id,
filename=filename_from_path(inner_path),
data=data,
display_path=target_path,
content_type=mimetypes.guess_type(inner_path)[0]
or "application/octet-stream",
metadata=metadata,
campaign_id=campaign_id,
conflict_strategy=conflict_strategy,
conflict_resolutions=conflict_resolutions,
is_admin=is_admin,
encryption_vault_id=encryption_vault_id,
)
)
stored_bytes += len(data)
if progress:
progress("storing", len(uploaded), total_files, stored_bytes, total_bytes)
del data
return uploaded
def _safe_member_path(value: str) -> str:
raw = str(value or "").replace("\\", "/").strip()
try:
oversized = (
len(raw) > ARCHIVE_MAX_PATH_BYTES
or len(raw.encode("utf-8")) > ARCHIVE_MAX_PATH_BYTES
)
except UnicodeEncodeError as exc:
raise FileStorageError("Archive member path is not valid Unicode") from exc
if oversized:
raise FileStorageError(
f"Archive member path exceeds {ARCHIVE_MAX_PATH_BYTES} UTF-8 bytes"
)
if not raw or "\x00" in raw or raw.startswith("/") or _WINDOWS_DRIVE_RE.match(raw):
raise FileStorageError(f"Unsafe archive member path {value!r}")
if any(part == ".." for part in raw.split("/")):
raise FileStorageError(f"Unsafe archive member path {value!r}")
try:
normalized = normalize_logical_path(raw.rstrip("/"))
except ValueError as exc:
raise FileStorageError(f"Unsafe archive member path {value!r}") from exc
if normalized.count("/") + 1 > ARCHIVE_MAX_PATH_DEPTH:
raise FileStorageError(
f"Archive member path exceeds {ARCHIVE_MAX_PATH_DEPTH} components"
)
return normalized
def _archive_size(archive_data: bytes | str | PathLike[str]) -> int:
if isinstance(archive_data, bytes):
return len(archive_data)
try:
return Path(archive_data).stat().st_size
except OSError as exc:
raise FileStorageError("Archive upload could not be read") from exc
def _archive_source(
archive_data: bytes | str | PathLike[str],
) -> BytesIO | str | PathLike[str]:
return BytesIO(archive_data) if isinstance(archive_data, bytes) else archive_data
def _open_tar(
archive_data: bytes | str | PathLike[str],
) -> tarfile.TarFile:
if isinstance(archive_data, bytes):
return tarfile.open(fileobj=BytesIO(archive_data), mode="r:*")
try:
return tarfile.open(name=archive_data, mode="r:*")
except OSError as exc:
raise FileStorageError("Archive upload could not be read") from exc