feat(campaign): enforce attachment delivery policies
This commit is contained in:
@@ -1,8 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import binascii
|
||||
from datetime import datetime
|
||||
import secrets
|
||||
import stat
|
||||
import struct
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
import zlib
|
||||
|
||||
try:
|
||||
import pyzipper
|
||||
@@ -10,6 +16,8 @@ except ImportError: # pragma: no cover
|
||||
pyzipper = None
|
||||
|
||||
ArchiveMember = tuple[Path, str]
|
||||
ZIP_METHOD_AES = "aes"
|
||||
ZIP_METHOD_STANDARD = "zip_standard"
|
||||
|
||||
|
||||
def _normalized_members(files: Iterable[Path | ArchiveMember]) -> list[ArchiveMember]:
|
||||
@@ -34,23 +42,17 @@ def create_zip_archive(
|
||||
output_path: Path,
|
||||
files: Iterable[Path | ArchiveMember],
|
||||
password: str = "",
|
||||
method: str = ZIP_METHOD_AES,
|
||||
) -> Path:
|
||||
"""Create a ZIP archive, using AES encryption when a password is supplied."""
|
||||
"""Create a ZIP archive, optionally using AES or legacy ZipCrypto encryption."""
|
||||
|
||||
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)
|
||||
if method == ZIP_METHOD_STANDARD:
|
||||
_create_zipcrypto_archive(output_path, members, password)
|
||||
return output_path
|
||||
_create_aes_archive(output_path, members, password)
|
||||
return output_path
|
||||
|
||||
with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as zip_file:
|
||||
@@ -59,7 +61,214 @@ def create_zip_archive(
|
||||
return output_path
|
||||
|
||||
|
||||
def create_encrypted_zip(output_path: Path, files: list[Path], password: str) -> Path:
|
||||
def create_encrypted_zip(output_path: Path, files: list[Path], password: str, method: str = ZIP_METHOD_AES) -> Path:
|
||||
"""Backward-compatible wrapper for the original per-rule ZIP helper."""
|
||||
|
||||
return create_zip_archive(output_path, files, password)
|
||||
return create_zip_archive(output_path, files, password, method)
|
||||
|
||||
|
||||
def _create_aes_archive(output_path: Path, members: list[ArchiveMember], password: str) -> None:
|
||||
if pyzipper is None:
|
||||
raise RuntimeError("pyzipper is required for writing AES 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)
|
||||
|
||||
|
||||
def _create_zipcrypto_archive(output_path: Path, members: list[ArchiveMember], password: str) -> None:
|
||||
entries: list[_CentralDirectoryEntry] = []
|
||||
password_bytes = password.encode("utf-8")
|
||||
with output_path.open("wb") as zip_file:
|
||||
for file_path, archive_name in members:
|
||||
local_header_offset = zip_file.tell()
|
||||
file_data = file_path.read_bytes()
|
||||
crc = binascii.crc32(file_data) & 0xFFFFFFFF
|
||||
compressed_data = _deflate(file_data)
|
||||
encrypted_data = _zipcrypto_encrypt(compressed_data, password_bytes, crc)
|
||||
compressed_size = len(encrypted_data)
|
||||
uncompressed_size = len(file_data)
|
||||
modified_at = datetime.fromtimestamp(file_path.stat().st_mtime)
|
||||
dos_time, dos_date = _dos_timestamp(modified_at)
|
||||
filename_bytes, flags = _filename_bytes(archive_name, encrypted=True)
|
||||
|
||||
zip_file.write(
|
||||
struct.pack(
|
||||
"<IHHHHHIIIHH",
|
||||
0x04034B50,
|
||||
20,
|
||||
flags,
|
||||
zipfile.ZIP_DEFLATED,
|
||||
dos_time,
|
||||
dos_date,
|
||||
crc,
|
||||
compressed_size,
|
||||
uncompressed_size,
|
||||
len(filename_bytes),
|
||||
0,
|
||||
)
|
||||
)
|
||||
zip_file.write(filename_bytes)
|
||||
zip_file.write(encrypted_data)
|
||||
entries.append(
|
||||
_CentralDirectoryEntry(
|
||||
filename_bytes=filename_bytes,
|
||||
flags=flags,
|
||||
dos_time=dos_time,
|
||||
dos_date=dos_date,
|
||||
crc=crc,
|
||||
compressed_size=compressed_size,
|
||||
uncompressed_size=uncompressed_size,
|
||||
external_attributes=_external_attributes(file_path),
|
||||
local_header_offset=local_header_offset,
|
||||
)
|
||||
)
|
||||
|
||||
central_directory_offset = zip_file.tell()
|
||||
for entry in entries:
|
||||
zip_file.write(
|
||||
struct.pack(
|
||||
"<IHHHHHHIIIHHHHHII",
|
||||
0x02014B50,
|
||||
20,
|
||||
20,
|
||||
entry.flags,
|
||||
zipfile.ZIP_DEFLATED,
|
||||
entry.dos_time,
|
||||
entry.dos_date,
|
||||
entry.crc,
|
||||
entry.compressed_size,
|
||||
entry.uncompressed_size,
|
||||
len(entry.filename_bytes),
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
entry.external_attributes,
|
||||
entry.local_header_offset,
|
||||
)
|
||||
)
|
||||
zip_file.write(entry.filename_bytes)
|
||||
|
||||
central_directory_size = zip_file.tell() - central_directory_offset
|
||||
zip_file.write(
|
||||
struct.pack(
|
||||
"<IHHHHIIH",
|
||||
0x06054B50,
|
||||
0,
|
||||
0,
|
||||
len(entries),
|
||||
len(entries),
|
||||
central_directory_size,
|
||||
central_directory_offset,
|
||||
0,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _CentralDirectoryEntry:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
filename_bytes: bytes,
|
||||
flags: int,
|
||||
dos_time: int,
|
||||
dos_date: int,
|
||||
crc: int,
|
||||
compressed_size: int,
|
||||
uncompressed_size: int,
|
||||
external_attributes: int,
|
||||
local_header_offset: int,
|
||||
) -> None:
|
||||
self.filename_bytes = filename_bytes
|
||||
self.flags = flags
|
||||
self.dos_time = dos_time
|
||||
self.dos_date = dos_date
|
||||
self.crc = crc
|
||||
self.compressed_size = compressed_size
|
||||
self.uncompressed_size = uncompressed_size
|
||||
self.external_attributes = external_attributes
|
||||
self.local_header_offset = local_header_offset
|
||||
|
||||
|
||||
def _deflate(data: bytes) -> bytes:
|
||||
compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
|
||||
return compressor.compress(data) + compressor.flush()
|
||||
|
||||
|
||||
def _zipcrypto_encrypt(data: bytes, password: bytes, crc: int) -> bytes:
|
||||
cipher = _ZipCrypto(password)
|
||||
header = secrets.token_bytes(11) + bytes([(crc >> 24) & 0xFF])
|
||||
return cipher.encrypt(header + data)
|
||||
|
||||
|
||||
class _ZipCrypto:
|
||||
def __init__(self, password: bytes) -> None:
|
||||
self.key0 = 0x12345678
|
||||
self.key1 = 0x23456789
|
||||
self.key2 = 0x34567890
|
||||
for value in password:
|
||||
self._update_keys(value)
|
||||
|
||||
def encrypt(self, data: bytes) -> bytes:
|
||||
encrypted = bytearray()
|
||||
for value in data:
|
||||
encrypted.append(value ^ self._decrypt_byte())
|
||||
self._update_keys(value)
|
||||
return bytes(encrypted)
|
||||
|
||||
def _decrypt_byte(self) -> int:
|
||||
temp = self.key2 | 2
|
||||
return ((temp * (temp ^ 1)) >> 8) & 0xFF
|
||||
|
||||
def _update_keys(self, value: int) -> None:
|
||||
self.key0 = _zipcrypto_crc32_byte(value, self.key0)
|
||||
self.key1 = (self.key1 + (self.key0 & 0xFF)) & 0xFFFFFFFF
|
||||
self.key1 = (self.key1 * 134775813 + 1) & 0xFFFFFFFF
|
||||
self.key2 = _zipcrypto_crc32_byte((self.key1 >> 24) & 0xFF, self.key2)
|
||||
|
||||
|
||||
def _build_zipcrypto_crc(value: int) -> int:
|
||||
for _ in range(8):
|
||||
if value & 1:
|
||||
value = 0xEDB88320 ^ (value >> 1)
|
||||
else:
|
||||
value >>= 1
|
||||
return value
|
||||
|
||||
|
||||
_ZIPCRYPTO_CRC_TABLE = [_build_zipcrypto_crc(value) for value in range(256)]
|
||||
|
||||
|
||||
def _zipcrypto_crc32_byte(value: int, crc: int) -> int:
|
||||
return ((crc >> 8) ^ _ZIPCRYPTO_CRC_TABLE[(crc ^ value) & 0xFF]) & 0xFFFFFFFF
|
||||
|
||||
|
||||
def _filename_bytes(filename: str, *, encrypted: bool) -> tuple[bytes, int]:
|
||||
flags = 0x1 if encrypted else 0
|
||||
try:
|
||||
return filename.encode("ascii"), flags
|
||||
except UnicodeEncodeError:
|
||||
return filename.encode("utf-8"), flags | 0x800
|
||||
|
||||
|
||||
def _dos_timestamp(value: datetime) -> tuple[int, int]:
|
||||
year = min(max(value.year, 1980), 2107)
|
||||
month = value.month if value.year == year else 1
|
||||
day = value.day if value.year == year else 1
|
||||
dos_time = value.hour << 11 | value.minute << 5 | value.second // 2
|
||||
dos_date = (year - 1980) << 9 | month << 5 | day
|
||||
return dos_time, dos_date
|
||||
|
||||
|
||||
def _external_attributes(path: Path) -> int:
|
||||
try:
|
||||
permissions = stat.S_IMODE(path.stat().st_mode)
|
||||
except OSError:
|
||||
return 0
|
||||
return permissions << 16
|
||||
|
||||
Reference in New Issue
Block a user