66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import zipfile
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
try:
|
|
import pyzipper
|
|
except ImportError: # pragma: no cover
|
|
pyzipper = None
|
|
|
|
ArchiveMember = tuple[Path, str]
|
|
|
|
|
|
def _normalized_members(files: Iterable[Path | ArchiveMember]) -> list[ArchiveMember]:
|
|
members: list[ArchiveMember] = []
|
|
used_names: set[str] = set()
|
|
for item in files:
|
|
path, requested_name = item if isinstance(item, tuple) else (item, item.name)
|
|
requested = Path(requested_name).name or path.name
|
|
stem = Path(requested).stem
|
|
suffix = Path(requested).suffix
|
|
candidate = requested
|
|
counter = 2
|
|
while candidate.casefold() in used_names:
|
|
candidate = f"{stem} ({counter}){suffix}"
|
|
counter += 1
|
|
used_names.add(candidate.casefold())
|
|
members.append((path, candidate))
|
|
return members
|
|
|
|
|
|
def create_zip_archive(
|
|
output_path: Path,
|
|
files: Iterable[Path | ArchiveMember],
|
|
password: str = "",
|
|
) -> Path:
|
|
"""Create a ZIP archive, using AES encryption when a password is supplied."""
|
|
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
members = _normalized_members(files)
|
|
if password:
|
|
if pyzipper is None:
|
|
raise RuntimeError("pyzipper is required for writing password-protected ZIP files")
|
|
with pyzipper.AESZipFile(
|
|
output_path,
|
|
"w",
|
|
compression=pyzipper.ZIP_DEFLATED,
|
|
encryption=pyzipper.WZ_AES,
|
|
) as zip_file:
|
|
zip_file.setpassword(password.encode("utf-8"))
|
|
for file_path, archive_name in members:
|
|
zip_file.write(file_path, arcname=archive_name)
|
|
return output_path
|
|
|
|
with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as zip_file:
|
|
for file_path, archive_name in members:
|
|
zip_file.write(file_path, arcname=archive_name)
|
|
return output_path
|
|
|
|
|
|
def create_encrypted_zip(output_path: Path, files: list[Path], password: str) -> Path:
|
|
"""Backward-compatible wrapper for the original per-rule ZIP helper."""
|
|
|
|
return create_zip_archive(output_path, files, password)
|