diff --git a/bookwyrm/connectors/abstract_connector.py b/bookwyrm/connectors/abstract_connector.py index 5ed57df1..0ab135b8 100644 --- a/bookwyrm/connectors/abstract_connector.py +++ b/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 resp + return None, None + + 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: diff --git a/bookwyrm/models/fields.py b/bookwyrm/models/fields.py index e61f912e..b506c11c 100644 --- a/bookwyrm/models/fields.py +++ b/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) - if not response: + image_content, extension = get_image(url) + 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] diff --git a/bookwyrm/tests/models/test_fields.py b/bookwyrm/tests/models/test_fields.py index 5bb7fecc..f7386c2e 100644 --- a/bookwyrm/tests/models/test_fields.py +++ b/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( - responses.GET, - "http://www.example.com/image.jpg", - body=image.tobytes(), - status=200, - ) + with open(image_file, "rb") as image_data: + responses.add( + responses.GET, + "http://www.example.com/image.jpg", + 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( - responses.GET, - "http://www.example.com/image.jpg", - body=image.tobytes(), - status=200, - ) + with open(image_file, "rb") as image_data: + responses.add( + responses.GET, + "http://www.example.com/image.jpg", + 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( - responses.GET, - "http://www.example.com/image.jpg", - body=image.tobytes(), - status=200, - ) + with open(image_file, "rb") as image_data: + responses.add( + responses.GET, + "http://www.example.com/image.jpg", + 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( - responses.GET, - "http://www.example.com/image.jpg", - body=another_image.tobytes(), - status=200, - ) + with open(another_image_file, "rb") as another_image: + responses.add( + responses.GET, + "http://www.example.com/image.jpg", + 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") diff --git a/bookwyrm/views/books/books.py b/bookwyrm/views/books/books.py index 7de2d0d2..e04230ba 100644 --- a/bookwyrm/views/books/books.py +++ b/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_content = ContentFile(image_file.content) + image_name = str(uuid4()) + "." + extension return [image_name, image_content] diff --git a/bw-dev b/bw-dev index 5248c887..11977ad0 100755 --- a/bw-dev +++ b/bw-dev @@ -209,6 +209,7 @@ case "$CMD" in echo " build" echo " clean" echo " black" + echo " prettier" echo " populate_streams [--stream=]" echo " populate_suggestions" echo " generate_thumbnails"