131 lines
5.3 KiB
Python
131 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
import mimetypes
|
|
import zipfile
|
|
from io import BytesIO
|
|
from os import PathLike
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from govoplan_files.backend.db.models import FileAsset
|
|
from govoplan_files.backend.storage.common import FileConflictResolution, FileStorageError, UploadedStoredFile
|
|
from govoplan_files.backend.storage.backends import StorageBackendError, get_storage_backend
|
|
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
|
|
|
|
|
|
_ZIP_READ_CHUNK_SIZE = 1024 * 1024
|
|
|
|
|
|
def _read_zip_member(
|
|
archive: zipfile.ZipFile,
|
|
info: zipfile.ZipInfo,
|
|
*,
|
|
max_file_bytes: int,
|
|
max_total_bytes: int,
|
|
current_total: int,
|
|
) -> tuple[bytes, int]:
|
|
parts: list[bytes] = []
|
|
actual_size = 0
|
|
with archive.open(info) as source:
|
|
while True:
|
|
read_size = min(_ZIP_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"ZIP member {info.filename!r} exceeds per-file limit")
|
|
if current_total + actual_size > max_total_bytes:
|
|
raise FileStorageError("ZIP is too large after extraction")
|
|
parts.append(chunk)
|
|
return b"".join(parts), current_total + actual_size
|
|
|
|
|
|
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 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,
|
|
max_files: int = 1000,
|
|
max_file_bytes: int = 50 * 1024 * 1024,
|
|
max_total_bytes: int = 250 * 1024 * 1024,
|
|
) -> list[UploadedStoredFile]:
|
|
uploaded: list[UploadedStoredFile] = []
|
|
total = 0
|
|
base_folder = normalize_folder(folder)
|
|
try:
|
|
source = BytesIO(zip_data) if isinstance(zip_data, bytes) else zip_data
|
|
with zipfile.ZipFile(source) as archive:
|
|
infos = [info for info in archive.infolist() if not info.is_dir()]
|
|
if len(infos) > max_files:
|
|
raise FileStorageError(f"ZIP contains too many files (limit {max_files})")
|
|
for info in infos:
|
|
if info.flag_bits & 0x1:
|
|
raise FileStorageError("Encrypted ZIP uploads are not supported")
|
|
if info.file_size < 0:
|
|
raise FileStorageError("Invalid ZIP member")
|
|
if info.file_size > max_file_bytes:
|
|
raise FileStorageError(f"ZIP member {info.filename!r} exceeds per-file limit")
|
|
if total + info.file_size > max_total_bytes:
|
|
raise FileStorageError("ZIP is too large after extraction")
|
|
inner_path = normalize_logical_path(info.filename)
|
|
target_path = f"{base_folder}/{inner_path}" if base_folder else inner_path
|
|
data, total = _read_zip_member(
|
|
archive,
|
|
info,
|
|
max_file_bytes=max_file_bytes,
|
|
max_total_bytes=max_total_bytes,
|
|
current_total=total,
|
|
)
|
|
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,
|
|
)
|
|
)
|
|
except zipfile.BadZipFile as exc:
|
|
raise FileStorageError("Invalid ZIP upload") from exc
|
|
return uploaded
|