# SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 Allaun """Compression Engine Module. Compression and archive utilities using standard library only. Key Features: - Data compression (gzip, bz2, lzma) - File compression and decompression - Archive creation and extraction (zip, tar, tar.gz, tar.bz2, tar.xz) - Path validation for security - Compression ratio estimation - Standard library only (no external dependencies) Dependencies: - gzip (standard library) - bz2 (standard library) - lzma (standard library) - tarfile (standard library) - zipfile (standard library) - pathlib (standard library) """ from __future__ import annotations import gzip import bz2 import lzma import tarfile import zipfile from pathlib import Path from typing import Optional, List, Union PathLike = Union[str, Path] class CompressionError(Exception): """Exception raised for compression and archive operations. Attributes: message: Explanation of the error """ pass class Compression: """Compression and archive operations handler. Provides static methods for compressing/decompressing data and files, as well as creating and extracting archives in various formats. Attributes: DATA_FORMATS: List of supported compression formats EXTENSION_MAP: Mapping of file extensions to format names FORMAT_EXTENSION: Mapping of format names to file extensions TAR_MODE_MAP: Mapping of tar format names to tarfile modes COMPRESSION_RATIOS: Estimated compression ratios by format and data type """ DATA_FORMATS = ['gzip', 'bz2', 'lzma'] EXTENSION_MAP = { '.gz': 'gzip', '.bz2': 'bz2', '.xz': 'lzma', '.zip': 'zip' } FORMAT_EXTENSION = { 'gzip': '.gz', 'bz2': '.bz2', 'lzma': '.xz', 'zip': '.zip' } TAR_MODE_MAP = { 'tar': 'w', 'tar.gz': 'w:gz', 'tar.bz2': 'w:bz2', 'tar.xz': 'w:xz' } COMPRESSION_RATIOS = { 'gzip': {'text': 0.3, 'binary': 0.6, 'image': 0.9, 'video': 0.95}, 'gz': {'text': 0.3, 'binary': 0.6, 'image': 0.9, 'video': 0.95}, 'bz2': {'text': 0.25, 'binary': 0.55, 'image': 0.9, 'video': 0.95}, 'lzma': {'text': 0.2, 'binary': 0.5, 'image': 0.9, 'video': 0.95}, 'xz': {'text': 0.2, 'binary': 0.5, 'image': 0.9, 'video': 0.95}, } @staticmethod def _ensure_path(path: PathLike) -> Path: """Ensure path is a Path object. Args: path: Path as string or Path object Returns: Path object """ return Path(path) @staticmethod def _validate_extraction_path(output_dir: Path, member_name: str) -> None: """Validate extraction path to prevent directory traversal attacks. Args: output_dir: Target extraction directory member_name: Name of archive member to extract Raises: CompressionError: If the extraction path is unsafe """ member_path = (output_dir / member_name).resolve() output_dir_resolved = output_dir.resolve() try: member_path.relative_to(output_dir_resolved) except ValueError: raise CompressionError( f"Unsafe extraction path: {member_name}" ) @staticmethod def _validate_format(format: str) -> None: """Validate that the format is supported. Args: format: Compression or archive format name Raises: CompressionError: If format is not supported """ if format not in Compression.DATA_FORMATS and \ format not in Compression.FORMAT_EXTENSION and \ format not in Compression.TAR_MODE_MAP: raise CompressionError(f"Unsupported format: {format}") @staticmethod def compress_data(data: bytes, format: str = 'gzip', level: int = 6) -> bytes: """Compress raw bytes data. Args: data: Bytes to compress format: Compression format ('gzip', 'bz2', or 'lzma') level: Compression level (1-9 for gzip/bz2, 0-9 for lzma) Returns: Compressed bytes Raises: CompressionError: If compression fails or format is unsupported """ Compression._validate_format(format) try: if format == 'gzip': return gzip.compress(data, compresslevel=level) if format == 'bz2': return bz2.compress(data, compresslevel=level) if format == 'lzma': return lzma.compress(data, preset=level) except Exception as e: raise CompressionError(f"Compression failed: {e}") from e raise CompressionError(f"Unsupported format: {format}") @staticmethod def decompress_data(data: bytes, format: str = 'gzip') -> bytes: """Decompress compressed bytes data. Args: data: Compressed bytes to decompress format: Compression format ('gzip', 'bz2', or 'lzma') Returns: Decompressed bytes Raises: CompressionError: If decompression fails or format is unsupported """ Compression._validate_format(format) try: if format == 'gzip': return gzip.decompress(data) if format == 'bz2': return bz2.decompress(data) if format == 'lzma': return lzma.decompress(data) except Exception as e: raise CompressionError(f"Decompression failed: {e}") from e raise CompressionError(f"Unsupported format: {format}") @staticmethod def compress_file( input_path: PathLike, output_path: Optional[PathLike] = None, format: str = 'gzip', remove_original: bool = False ) -> Path: """Compress a file. Args: input_path: Path to the input file output_path: Path for the output compressed file (None = auto-generate) format: Compression format ('gzip', 'bz2', 'lzma', or 'zip') remove_original: Whether to remove the original file after compression Returns: Path to the compressed file Raises: CompressionError: If compression fails """ input_path = Compression._ensure_path(input_path) if not input_path.exists(): raise CompressionError("Input file does not exist") Compression._validate_format(format) if output_path is None: ext = Compression.FORMAT_EXTENSION.get(format, f'.{format}') output_path = input_path.with_suffix(input_path.suffix + ext) else: output_path = Compression._ensure_path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) try: if format in Compression.DATA_FORMATS: with open(input_path, 'rb') as src: data = src.read() compressed = Compression.compress_data(data, format) with open(output_path, 'wb') as dst: dst.write(compressed) elif format == 'zip': with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zf: zf.write(input_path, arcname=input_path.name) else: raise CompressionError(f"Unsupported format: {format}") except Exception as e: raise CompressionError(f"File compression failed: {e}") from e if remove_original: try: input_path.unlink() except Exception as e: raise CompressionError(f"Failed to remove original: {e}") from e return output_path @staticmethod def create_archive( files: List[PathLike], output_path: PathLike, format: str = 'zip' ) -> Path: """Create an archive from a list of files. Args: files: List of paths to files to include output_path: Path to the resulting archive format: Archive format ('zip', 'tar', 'tar.gz', etc.) Returns: Path to the created archive Raises: CompressionError: If archive creation fails """ output_path = Compression._ensure_path(output_path) Compression._validate_format(format) output_path.parent.mkdir(parents=True, exist_ok=True) try: if format == 'zip': with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zf: for f in files: f = Compression._ensure_path(f) if f.exists(): zf.write(f, arcname=f.name) else: raise CompressionError(f"File not found: {f}") elif format in Compression.TAR_MODE_MAP: mode = Compression.TAR_MODE_MAP[format] with tarfile.open(output_path, mode) as tf: for f in files: f = Compression._ensure_path(f) if f.exists(): tf.add(f, arcname=f.name) else: raise CompressionError(f"File not found: {f}") else: raise CompressionError(f"Unsupported format: {format}") return output_path except Exception as e: if isinstance(e, CompressionError): raise raise CompressionError(f"Archive creation failed: {e}") from e @staticmethod def decompress_file( input_path: PathLike, output_path: Optional[PathLike] = None, format: Optional[str] = None, remove_compressed: bool = False ) -> Path: """Decompress a compressed file. Args: input_path: Path to the compressed file output_path: Path for the output decompressed file (None = auto-generate) format: Compression format (None = auto-detect from extension) remove_compressed: Whether to remove the compressed file after decompression Returns: Path to the decompressed file Raises: CompressionError: If decompression fails """ input_path = Compression._ensure_path(input_path) if not input_path.exists(): raise CompressionError("Input file does not exist") if format is None: format = Compression.EXTENSION_MAP.get(input_path.suffix.lower()) if not format: raise CompressionError("Cannot auto-detect format") Compression._validate_format(format) if output_path is None: output_path = input_path.with_suffix('') else: output_path = Compression._ensure_path(output_path) try: if format in Compression.DATA_FORMATS: with open(input_path, 'rb') as src: data = src.read() decompressed = Compression.decompress_data(data, format) with open(output_path, 'wb') as dst: dst.write(decompressed) elif format == 'zip': with zipfile.ZipFile(input_path, 'r') as zf: for name in zf.namelist(): Compression._validate_extraction_path(output_path.parent, name) zf.extractall(output_path.parent) else: raise CompressionError(f"Unsupported format: {format}") except Exception as e: raise CompressionError(f"File decompression failed: {e}") from e if remove_compressed: input_path.unlink() return output_path @staticmethod def extract_archive( archive_path: PathLike, output_dir: PathLike, format: Optional[str] = None, PASSWORD_REMOVED: Optional[bytes] = None ) -> List[Path]: """Extract archive (zip, tar, tar.gz, tar.bz2, tar.xz) to output directory. Args: archive_path: Path to the archive file output_dir: Directory to extract to format: Optional format hint ('zip', 'tar', 'tar.gz', etc.) PASSWORD_REMOVED: Optional PASSWORD_REMOVED for encrypted archives Returns: List of extracted file paths Raises: CompressionError: If archive cannot be extracted """ archive_path = Compression._ensure_path(archive_path) output_dir = Compression._ensure_path(output_dir) if not archive_path.exists(): raise CompressionError("Archive not found") # Auto-detect format if not provided detected = format if detected is None: suffix = archive_path.suffix.lower() if suffix == '.zip': detected = 'zip' elif archive_path.name.endswith('.tar.gz'): detected = 'tar.gz' elif archive_path.name.endswith('.tar.bz2'): detected = 'tar.bz2' elif archive_path.name.endswith('.tar.xz'): detected = 'tar.xz' elif suffix == '.tar': detected = 'tar' else: raise CompressionError(f"Cannot auto-detect format for: {archive_path}") output_dir.mkdir(parents=True, exist_ok=True) extracted_files: List[Path] = [] try: if detected == 'zip': with zipfile.ZipFile(archive_path, 'r') as zf: if PASSWORD_REMOVED: zf.setPASSWORD_REMOVED(PASSWORD_REMOVED) # Validate all paths before extracting for name in zf.namelist(): Compression._validate_extraction_path(output_dir, name) zf.extractall(output_dir) extracted_files = [output_dir / name for name in zf.namelist()] elif detected in Compression.TAR_MODE_MAP: with tarfile.open(archive_path, 'r:*') as tf: # Validate all paths before extracting members = [m for m in tf.getmembers() if m.name not in ('.', './', '..', '../')] for member in members: Compression._validate_extraction_path(output_dir, member.name) tf.extractall(output_dir, members=members) extracted_files = [output_dir / member.name for member in members] else: raise CompressionError(f"Unsupported format: {detected}") return extracted_files except Exception as e: if isinstance(e, CompressionError): raise raise CompressionError(f"Archive extraction failed: {e}") from e @staticmethod def get_compression_ratio(original: int, compressed: int) -> float: """Calculate compression ratio. Args: original: Original (uncompressed) size in bytes compressed: Compressed size in bytes Returns: Compression ratio (original / compressed), or 0 if compressed is 0 """ return original / compressed if compressed else 0.0 @staticmethod def estimate_compressed_size( size: int, format: str = 'gzip', data_type: str = 'text' ) -> int: """Estimate compressed size based on historical ratios. Args: size: Original size in bytes format: Compression format data_type: Type of data ('text', 'binary', 'image', or 'video') Returns: Estimated compressed size in bytes """ format_key = format.replace('tar.', '') ratios = Compression.COMPRESSION_RATIOS.get(format_key) if ratios: ratio = ratios.get(data_type, 0.5) else: ratio = 0.5 return int(size * ratio)