Last active
December 6, 2024 10:33
-
-
Save archatas/a463ce624cbd587431884dfbab322ba2 to your computer and use it in GitHub Desktop.
Check the media directory for image decompression bombs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import os | |
| import imghdr | |
| import warnings | |
| from PIL import Image, ImageFile | |
| def configure_image_safety(): | |
| """ | |
| Configure image processing safety settings to prevent decompression bomb issues. | |
| """ | |
| # Increase the maximum allowed pixels to prevent warnings | |
| Image.MAX_IMAGE_PIXELS = 1000000000 # 1 billion pixels (adjust as needed) | |
| # Correctly filter warnings | |
| warnings.simplefilter('ignore', Image.DecompressionBombWarning) | |
| warnings.simplefilter('ignore', Image.DecompressionBombError) | |
| # Configure PIL to raise exceptions for large images | |
| ImageFile.LOAD_TRUNCATED_IMAGES = False | |
| def is_decompression_bomb(file_path, max_pixels=100_000_000): | |
| """ | |
| Check if an image is a potential decompression bomb. | |
| Args: | |
| - file_path (str): Path to the image file | |
| - max_pixels (int): Maximum number of pixels allowed before considering it a bomb | |
| Returns: | |
| - bool: True if image is a potential decompression bomb, False otherwise | |
| """ | |
| try: | |
| # Check if it's actually an image | |
| if not imghdr.what(file_path): | |
| return False | |
| # Use a try-except block to handle potential image processing issues | |
| with Image.open(file_path) as img: | |
| # Calculate total pixels | |
| pixels = img.width * img.height | |
| # Check if pixels exceed threshold | |
| return pixels > max_pixels | |
| except Exception as e: | |
| print(f"Error processing {file_path}: {e}") | |
| return False | |
| def find_image_bombs(directory, max_pixels=100_000_000): | |
| """ | |
| Recursively find image decompression bombs in a directory. | |
| Args: | |
| - directory (str): Root directory to search | |
| - max_pixels (int): Maximum number of pixels allowed | |
| Returns: | |
| - list: Paths of potential decompression bomb images | |
| """ | |
| # Configure image safety settings | |
| configure_image_safety() | |
| bombs = [] | |
| # Walk through directory recursively | |
| for root, _, files in os.walk(directory): | |
| for file in files: | |
| file_path = os.path.join(root, file) | |
| try: | |
| # Check if file is an image and a potential bomb | |
| if is_decompression_bomb(file_path, max_pixels): | |
| bombs.append(file_path) | |
| except Exception as e: | |
| print(f"Error processing {file_path}: {e}") | |
| return bombs | |
| # Alternative warning suppression method | |
| def suppress_decompression_bomb_warnings(): | |
| """ | |
| Alternative method to suppress decompression bomb warnings. | |
| """ | |
| import warnings | |
| def warn(*args, **kwargs): | |
| pass | |
| warnings.warn = warn | |
| # Example usage | |
| if __name__ == '__main__': | |
| # Option 1: Use configure_image_safety() | |
| directory = '/path/to/your/images' | |
| bombs = find_image_bombs(directory, max_pixels=50_000_000) | |
| # Option 2: Alternative warning suppression | |
| # suppress_decompression_bomb_warnings() | |
| print("Potential Decompression Bombs:") | |
| for bomb in bombs: | |
| print(bomb) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment