Merge pull request #6 from hugovk/bomb2

Change the decompression bomb warning into a subclass of RuntimeWarning to aid filtering

PIL/Image.py

@@ -31,6 +31,9 @@ from PIL import VERSION, PILLOW_VERSION, _plugins
 import warnings
 
 
+class DecompressionBombWarning(RuntimeWarning):
+    pass
+
 class _imaging_not_installed:
     # module placeholder
     def __getattr__(self, id):
@@ -2187,7 +2190,7 @@ def _decompression_bomb_check(size):
             "Image size (%d pixels) exceeds limit of %d pixels, "
             "could be decompression bomb DOS attack." %
             (pixels, MAX_IMAGE_PIXELS),
-            RuntimeWarning)
+            DecompressionBombWarning)
 
 
 def open(fp, mode="r"):

Tests/test_decompression_bomb.py

@@ -1,37 +1,45 @@
-from tester import *
+from helper import unittest, PillowTestCase, tearDownModule
 
 from PIL import Image
 
 test_file = "Images/lena.ppm"
 
-
-def test_no_warning_small_file():
-    # Implicit assert: no warning.
-    # A warning would cause a failure.
-    Image.open(test_file)
+ORIGINAL_LIMIT = Image.MAX_IMAGE_PIXELS
 
 
-def test_no_warning_no_limit():
-    # Arrange
-    # Turn limit off
-    Image.MAX_IMAGE_PIXELS = None
-    assert_equal(Image.MAX_IMAGE_PIXELS, None)
+class TestDecompressionBomb(PillowTestCase):
 
-    # Act / Assert
-    # Implicit assert: no warning.
-    # A warning would cause a failure.
-    Image.open(test_file)
+    def tearDown(self):
+        Image.MAX_IMAGE_PIXELS = ORIGINAL_LIMIT
 
+    def test_no_warning_small_file(self):
+        # Implicit assert: no warning.
+        # A warning would cause a failure.
+        Image.open(test_file)
 
-def test_warning():
-    # Arrange
-    # Set limit to a low, easily testable value
-    Image.MAX_IMAGE_PIXELS = 10
-    assert_equal(Image.MAX_IMAGE_PIXELS, 10)
+    def test_no_warning_no_limit(self):
+        # Arrange
+        # Turn limit off
+        Image.MAX_IMAGE_PIXELS = None
+        self.assertEqual(Image.MAX_IMAGE_PIXELS, None)
 
-    # Act / Assert
-    assert_warning(
-        RuntimeWarning,
-        lambda: Image.open(test_file))
+        # Act / Assert
+        # Implicit assert: no warning.
+        # A warning would cause a failure.
+        Image.open(test_file)
+
+    def test_warning(self):
+        # Arrange
+        # Set limit to a low, easily testable value
+        Image.MAX_IMAGE_PIXELS = 10
+        self.assertEqual(Image.MAX_IMAGE_PIXELS, 10)
+
+        # Act / Assert
+        self.assert_warning(
+            Image.DecompressionBombWarning,
+            lambda: Image.open(test_file))
+
+if __name__ == '__main__':
+    unittest.main()
 
 # End of file