Merge pull request #1897 from bookwyrm-social/verify-image-types

Check image extensions before saving

bookwyrm/connectors/abstract_connector.py

@@ -1,7 +1,9 @@
 """ functionality outline for a book data connector """
 from abc import ABC, abstractmethod
+import imghdr
 import logging
 
+from django.core.files.base import ContentFile
 from django.db import transaction
 import requests
 from requests.exceptions import RequestException
@@ -290,10 +292,18 @@ def get_image(url, timeout=10):
         )
     except RequestException as err:
         logger.exception(err)
-        return None
+        return None, None
+
     if not resp.ok:
-        return None
+        return None, None
-    return resp
+
+    image_content = ContentFile(resp.content)
+    extension = imghdr.what(None, image_content.read())
+    if not extension:
+        logger.exception("File requested was not an image: %s", url)
+        return None, None
+
+    return image_content, extension
 
 
 class Mapping:

bookwyrm/models/fields.py

@@ -1,6 +1,5 @@
 """ activitypub-aware django model fields """
 from dataclasses import MISSING
-import imghdr
 import re
 from uuid import uuid4
 from urllib.parse import urljoin
@@ -9,7 +8,6 @@ import dateutil.parser
 from dateutil.parser import ParserError
 from django.contrib.postgres.fields import ArrayField as DjangoArrayField
 from django.core.exceptions import ValidationError
-from django.core.files.base import ContentFile
 from django.db import models
 from django.forms import ClearableFileInput, ImageField as DjangoImageField
 from django.utils import timezone
@@ -443,12 +441,10 @@ class ImageField(ActivitypubFieldMixin, models.ImageField):
         except ValidationError:
             return None
 
-        response = get_image(url)
+        image_content, extension = get_image(url)
-        if not response:
+        if not image_content:
             return None
 
-        image_content = ContentFile(response.content)
-        extension = imghdr.what(None, image_content.read()) or ""
         image_name = f"{uuid4()}.{extension}"
         return [image_name, image_content]
 

bookwyrm/tests/models/test_fields.py

@@ -443,18 +443,17 @@ class ModelFields(TestCase):
         image_file = pathlib.Path(__file__).parent.joinpath(
             "../../static/images/default_avi.jpg"
         )
-        image = Image.open(image_file)
-        output = BytesIO()
-        image.save(output, format=image.format)
-
         instance = fields.ImageField()
 
-        responses.add(
+        with open(image_file, "rb") as image_data:
-            responses.GET,
+            responses.add(
-            "http://www.example.com/image.jpg",
+                responses.GET,
-            body=image.tobytes(),
+                "http://www.example.com/image.jpg",
-            status=200,
+                body=image_data.read(),
-        )
+                status=200,
+                content_type="image/jpeg",
+                stream=True,
+            )
         loaded_image = instance.field_from_activity("http://www.example.com/image.jpg")
         self.assertIsInstance(loaded_image, list)
         self.assertIsInstance(loaded_image[1], ContentFile)
@@ -465,18 +464,18 @@ class ModelFields(TestCase):
         image_file = pathlib.Path(__file__).parent.joinpath(
             "../../static/images/default_avi.jpg"
         )
-        image = Image.open(image_file)
-        output = BytesIO()
-        image.save(output, format=image.format)
 
         instance = fields.ImageField(activitypub_field="cover", name="cover")
 
-        responses.add(
+        with open(image_file, "rb") as image_data:
-            responses.GET,
+            responses.add(
-            "http://www.example.com/image.jpg",
+                responses.GET,
-            body=image.tobytes(),
+                "http://www.example.com/image.jpg",
-            status=200,
+                body=image_data.read(),
-        )
+                content_type="image/jpeg",
+                status=200,
+                stream=True,
+            )
         book = Edition.objects.create(title="hello")
 
         MockActivity = namedtuple("MockActivity", ("cover"))
@@ -491,18 +490,18 @@ class ModelFields(TestCase):
         image_file = pathlib.Path(__file__).parent.joinpath(
             "../../static/images/default_avi.jpg"
         )
-        image = Image.open(image_file)
-        output = BytesIO()
-        image.save(output, format=image.format)
 
         instance = fields.ImageField(activitypub_field="cover", name="cover")
 
-        responses.add(
+        with open(image_file, "rb") as image_data:
-            responses.GET,
+            responses.add(
-            "http://www.example.com/image.jpg",
+                responses.GET,
-            body=image.tobytes(),
+                "http://www.example.com/image.jpg",
-            status=200,
+                body=image_data.read(),
-        )
+                status=200,
+                content_type="image/jpeg",
+                stream=True,
+            )
         book = Edition.objects.create(title="hello")
 
         MockActivity = namedtuple("MockActivity", ("cover"))
@@ -565,18 +564,18 @@ class ModelFields(TestCase):
         another_image_file = pathlib.Path(__file__).parent.joinpath(
             "../../static/images/logo.png"
         )
-        another_image = Image.open(another_image_file)
-        another_output = BytesIO()
-        another_image.save(another_output, format=another_image.format)
 
         instance = fields.ImageField(activitypub_field="cover", name="cover")
 
-        responses.add(
+        with open(another_image_file, "rb") as another_image:
-            responses.GET,
+            responses.add(
-            "http://www.example.com/image.jpg",
+                responses.GET,
-            body=another_image.tobytes(),
+                "http://www.example.com/image.jpg",
-            status=200,
+                body=another_image.read(),
-        )
+                status=200,
+                content_type="image/jpeg",
+                stream=True,
+            )
 
         MockActivity = namedtuple("MockActivity", ("cover"))
         mock_activity = MockActivity("http://www.example.com/image.jpg")

bookwyrm/views/books/books.py

@@ -2,7 +2,6 @@
 from uuid import uuid4
 
 from django.contrib.auth.decorators import login_required, permission_required
-from django.core.files.base import ContentFile
 from django.core.paginator import Paginator
 from django.db.models import Avg, Q
 from django.http import Http404
@@ -144,13 +143,12 @@ def upload_cover(request, book_id):
 def set_cover_from_url(url):
     """load it from a url"""
     try:
-        image_file = get_image(url)
+        image_content, extension = get_image(url)
     except:  # pylint: disable=bare-except
         return None
-    if not image_file:
+    if not image_content:
         return None
-    image_name = str(uuid4()) + "." + url.split(".")[-1]
+    image_name = str(uuid4()) + "." + extension
-    image_content = ContentFile(image_file.content)
     return [image_name, image_content]
 
 

bw-dev

@@ -209,6 +209,7 @@ case "$CMD" in
         echo "    build"
         echo "    clean"
         echo "    black"
+        echo "    prettier"
         echo "    populate_streams [--stream=<stream name>]"
         echo "    populate_suggestions"
         echo "    generate_thumbnails"