Use Pillow for image processing

Fixes #207.

flaskbb/fixtures/settings.py

@@ -18,7 +18,7 @@ def available_themes():
 
 
 def available_avatar_types():
-    return [("image/png", "PNG"), ("image/jpeg", "JPG"), ("image/gif", "GIF")]
+    return [("PNG", "PNG"), ("JPEG", "JPG"), ("GIF", "GIF")]
 
 
 def available_languages():
@@ -185,7 +185,7 @@ fixture = (
                 'description':  "The allowed size of the avatar in kilobytes."
             }),
             ('avatar_types', {
-                'value':        ["image/png", "image/jpeg", "image/gif"],
+                'value':        ["PNG", "JPEG", "GIF"],
                 'value_type':   "selectmultiple",
                 'extra':        {"choices": available_avatar_types},
                 'name':         "Avatar Types",

flaskbb/utils/helpers.py

@@ -12,10 +12,9 @@ import re
 import time
 import itertools
 import operator
-import struct
-from io import BytesIO
 from datetime import datetime, timedelta
 from pytz import UTC
+from PIL import ImageFile
 
 import requests
 import unidecode
@@ -415,82 +414,42 @@ def format_quote(username, content):
 
 
 def get_image_info(url):
-    """Returns the content-type, image size (kb), height and width of a image
-    without fully downloading it. It will just download the first 1024 bytes.
+    """Returns the content-type, image size (kb), height and width of
+    an image without fully downloading it.
 
-    LICENSE: New BSD License (taken from the start page of the repository)
-    https://code.google.com/p/bfg-pages/source/browse/trunk/pages/getimageinfo.py
+    :param url: The URL of the image.
     """
     r = requests.get(url, stream=True)
     image_size = r.headers.get("content-length")
     image_size = float(image_size) / 1000  # in kilobyte
+    image_max_size = 10000
+    image_data = {
+        "content_type": "",
+        "size": image_size,
+        "width": 0,
+        "height": 0
+    }
 
-    data = r.raw.read(1024)
-    size = len(data)
-    height = -1
-    width = -1
-    content_type = ''
-
-    if size:
-        size = int(size)
-
-    # handle GIFs
-    if (size >= 10) and data[:6] in (b'GIF87a', b'GIF89a'):
-        # Check to see if content_type is correct
-        content_type = 'image/gif'
-        w, h = struct.unpack(b'<HH', data[6:10])
-        width = int(w)
-        height = int(h)
-
-    # See PNG 2. Edition spec (http://www.w3.org/TR/PNG/)
-    # Bytes 0-7 are below, 4-byte chunk length, then 'IHDR'
-    # and finally the 4-byte width, height
-    elif ((size >= 24) and data.startswith(b'\211PNG\r\n\032\n') and
-            (data[12:16] == b'IHDR')):
-        content_type = 'image/png'
-        w, h = struct.unpack(b">LL", data[16:24])
-        width = int(w)
-        height = int(h)
-
-    # Maybe this is for an older PNG version.
-    elif (size >= 16) and data.startswith(b'\211PNG\r\n\032\n'):
-        # Check to see if we have the right content type
-        content_type = 'image/png'
-        w, h = struct.unpack(b">LL", data[8:16])
-        width = int(w)
-        height = int(h)
-
-    # handle JPEGs
-    elif (size >= 2) and data.startswith(b'\377\330'):
-        content_type = 'image/jpeg'
-        jpeg = BytesIO(data)
-        jpeg.read(2)
-        b = jpeg.read(1)
-        try:
-            while (b and ord(b) != 0xDA):
-
-                while (ord(b) != 0xFF):
-                    b = jpeg.read(1)
-
-                while (ord(b) == 0xFF):
-                    b = jpeg.read(1)
-
-                if (ord(b) >= 0xC0 and ord(b) <= 0xC3):
-                    jpeg.read(3)
-                    h, w = struct.unpack(b">HH", jpeg.read(4))
-                    break
-                else:
-                    jpeg.read(int(struct.unpack(b">H", jpeg.read(2))[0]) - 2)
-                b = jpeg.read(1)
-            width = int(w)
-            height = int(h)
-        except struct.error:
-            pass
-        except ValueError:
-            pass
-
-    return {"content-type": content_type, "size": image_size,
-            "width": width, "height": height}
+    # lets set a hard limit of 10MB
+    if image_size > image_max_size:
+        return image_data
+
+    data = None
+    parser = ImageFile.Parser()
+
+    while True:
+        data = r.raw.read(1024)
+        if not data:
+            break
+
+        parser.feed(data)
+        if parser.image:
+            image_data["content_type"] = parser.image.format
+            image_data["width"] = parser.image.size[0]
+            image_data["height"] = parser.image.size[1]
+            break
+
+    return image_data
 
 
 def check_image(url):
@@ -506,8 +465,15 @@ def check_image(url):
     img_info = get_image_info(url)
     error = None
 
-    if not img_info["content-type"] in flaskbb_config["AVATAR_TYPES"]:
-        error = "Image type is not allowed. Allowed types are: {}".format(
+    if img_info["size"] > flaskbb_config["AVATAR_SIZE"]:
+        error = "Image is too big! {}kb are allowed.".format(
+            flaskbb_config["AVATAR_SIZE"]
+        )
+        return error, False
+
+    if not img_info["content_type"] in flaskbb_config["AVATAR_TYPES"]:
+        error = "Image type {} is not allowed. Allowed types are: {}".format(
+            img_info["content_type"],
             ", ".join(flaskbb_config["AVATAR_TYPES"])
         )
         return error, False
@@ -524,10 +490,4 @@ def check_image(url):
         )
         return error, False
 
-    if img_info["size"] > flaskbb_config["AVATAR_SIZE"]:
-        error = "Image is too big! {}kb are allowed.".format(
-            flaskbb_config["AVATAR_SIZE"]
-        )
-        return error, False
-
     return error, True

requirements.txt

@@ -29,6 +29,7 @@ limits==1.1.1
 Mako==1.0.4
 MarkupSafe==0.23
 mistune==0.7.3
+Pillow==3.3.1
 Pygments==2.1.3
 python-editor==1.0.1
 pytz==2016.6.1

setup.py

@@ -93,6 +93,7 @@ setup(
         'Mako',
         'MarkupSafe',
         'mistune',
+        'Pillow',
         'Pygments',
         'python-editor',
         'pytz',

tests/unit/utils/test_helpers.py

@@ -118,31 +118,36 @@ def test_get_image_info():
     gif = "http://i.imgur.com/l3Vmp4m.gif"
     png = "http://i.imgur.com/JXzKxNs.png"
 
+    # Issue #207 Image - This works now
+    #issue_img = "http://b.reich.io/gtlbjc.jpg"
+    #issue_img = get_image_info(issue_img)
+    #assert issue_img["content_type"] == "JPEG"
+
     jpg_img = get_image_info(jpg)
-    assert jpg_img["content-type"] == "image/jpeg"
+    assert jpg_img["content_type"] == "JPEG"
     assert jpg_img["height"] == 1024
     assert jpg_img["width"] == 1280
     assert jpg_img["size"] == 209.06
 
     gif_img = get_image_info(gif)
-    assert gif_img["content-type"] == "image/gif"
+    assert gif_img["content_type"] == "GIF"
     assert gif_img["height"] == 168
     assert gif_img["width"] == 400
     assert gif_img["size"] == 576.138
 
     png_img = get_image_info(png)
-    assert png_img["content-type"] == "image/png"
+    assert png_img["content_type"] == "PNG"
     assert png_img["height"] == 1080
     assert png_img["width"] == 1920
     assert png_img["size"] == 269.409
 
 
 def test_check_image(default_settings):
-    # test200x100.png
+    # test200_100.png
     img_width = "http://i.imgur.com/4dAWAZI.png"
-    # test100x200.png
+    # test100_200.png
     img_height = "http://i.imgur.com/I7GwF3D.png"
-    # test100x100.png
+    # test100_100.png
     img_ok = "http://i.imgur.com/CYV6NzT.png"
     # random too big image
     img_size = "http://i.imgur.com/l3Vmp4m.gif"