From 8df3009cb2ee25abb20371f685251d402cb700aa Mon Sep 17 00:00:00 2001
From: Nathan Beals <850796+ndbeals@users.noreply.github.com>
Date: Sat, 23 Apr 2022 00:11:04 -0400
Subject: [PATCH] Update image upload handler to be content-type aware Update
 handle_image: made filetype required and not optional Updated handle_image
 usage to reflect changes

---
 cookbook/helper/image_processing.py | 10 ++++++----
 cookbook/integration/integration.py |  2 +-
 cookbook/views/api.py               |  7 +++++--
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/cookbook/helper/image_processing.py b/cookbook/helper/image_processing.py
index d9a334acb..552bf1312 100644
--- a/cookbook/helper/image_processing.py
+++ b/cookbook/helper/image_processing.py
@@ -38,10 +38,12 @@ def get_filetype(name):
 
 # TODO this whole file needs proper documentation, refactoring, and testing
 # TODO also add env variable to define which images sizes should be compressed
-def handle_image(request, image_object, filetype='.jpeg'):
+# filetype argument can not be optional, otherwise this function will treat all images as if they were a jpeg
+# Because it's no longer optional, no reason to return it
+def handle_image(request, image_object, filetype): 
     if (image_object.size / 1000) > 500:  # if larger than 500 kb compress
         if filetype == '.jpeg' or filetype == '.jpg':
-            return rescale_image_jpeg(image_object), filetype
+            return rescale_image_jpeg(image_object)
         if filetype == '.png':
-            return rescale_image_png(image_object), filetype
-    return image_object, filetype
+            return rescale_image_png(image_object)
+    return image_object
diff --git a/cookbook/integration/integration.py b/cookbook/integration/integration.py
index 6493c266d..c3dc9dfb7 100644
--- a/cookbook/integration/integration.py
+++ b/cookbook/integration/integration.py
@@ -253,7 +253,7 @@ class Integration:
         :param image_file: ByteIO stream containing the image
         :param filetype: type of file to write bytes to, default to .jpeg if unknown
         """
-        recipe.image = File(handle_image(self.request, File(image_file, name='image'), filetype=filetype)[0], name=f'{uuid.uuid4()}_{recipe.pk}{filetype}')
+        recipe.image = File(handle_image(self.request, File(image_file, name='image'), filetype=filetype), name=f'{uuid.uuid4()}_{recipe.pk}{filetype}')
         recipe.save()
 
     def get_recipe_from_file(self, file):
diff --git a/cookbook/views/api.py b/cookbook/views/api.py
index 15e3cf299..635ee15fa 100644
--- a/cookbook/views/api.py
+++ b/cookbook/views/api.py
@@ -1,5 +1,6 @@
 import io
 import json
+import mimetypes
 import re
 import uuid
 from collections import OrderedDict
@@ -772,14 +773,16 @@ class RecipeViewSet(viewsets.ModelViewSet):
         if serializer.is_valid():
             serializer.save()
             image = None
+            filetype = ".jpeg" # fall-back to .jpeg, even if wrong, at least users will know it's an image and most image viewers can open it correctly anyways
 
             if 'image' in serializer.validated_data:
                 image = obj.image
+                filetype = mimetypes.guess_extension(serializer.validated_data['image'].content_type) or filetype
             elif 'image_url' in serializer.validated_data:
                 try:
                     response = requests.get(serializer.validated_data['image_url'])
                     image = File(io.BytesIO(response.content))
-                    print('test')
+                    filetype = mimetypes.guess_extension(response.headers['content-type']) or filetype
                 except UnidentifiedImageError as e:
                     print(e)
                     pass
@@ -791,7 +794,7 @@ class RecipeViewSet(viewsets.ModelViewSet):
                     pass
 
             if image is not None:
-                img, filetype = handle_image(request, image)
+                img = handle_image(request, image, filetype)
                 obj.image = File(img, name=f'{uuid.uuid4()}_{obj.pk}{filetype}')
                 obj.save()
                 return Response(serializer.data)