76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import mimetypes
|
|
import zipfile
|
|
from io import BytesIO
|
|
from typing import Iterable
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.db.models import FileAsset
|
|
from app.storage.common import FileStorageError, UploadedStoredFile
|
|
from app.storage.files import create_file_asset, read_asset_bytes
|
|
from app.storage.paths import filename_from_path, normalize_folder, normalize_logical_path
|
|
|
|
|
|
def create_zip_bytes(session: Session, assets: Iterable[FileAsset]) -> bytes:
|
|
buffer = BytesIO()
|
|
with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as archive:
|
|
for asset in assets:
|
|
data, _, _ = read_asset_bytes(session, asset)
|
|
archive.writestr(asset.display_path, data)
|
|
buffer.seek(0)
|
|
return buffer.getvalue()
|
|
|
|
|
|
def extract_zip_upload(
|
|
session: Session,
|
|
*,
|
|
tenant_id: str,
|
|
owner_type: str,
|
|
owner_id: str,
|
|
user_id: str,
|
|
zip_data: bytes,
|
|
folder: str | None,
|
|
campaign_id: str | None,
|
|
conflict_strategy: str = "reject",
|
|
conflict_resolutions: Iterable[FileConflictResolution] | None = None,
|
|
is_admin: bool = False,
|
|
max_files: int = 1000,
|
|
max_total_bytes: int = 250 * 1024 * 1024,
|
|
) -> list[UploadedStoredFile]:
|
|
uploaded: list[UploadedStoredFile] = []
|
|
total = 0
|
|
base_folder = normalize_folder(folder)
|
|
with zipfile.ZipFile(BytesIO(zip_data)) 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.file_size < 0:
|
|
raise FileStorageError("Invalid ZIP member")
|
|
total += info.file_size
|
|
if total > 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 = archive.read(info)
|
|
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",
|
|
campaign_id=campaign_id,
|
|
conflict_strategy=conflict_strategy,
|
|
conflict_resolutions=conflict_resolutions,
|
|
is_admin=is_admin,
|
|
)
|
|
)
|
|
return uploaded
|