Skip to content

Commit

Permalink
fixes #57: add magic_header check
Browse files Browse the repository at this point in the history
  • Loading branch information
Maxence Guindon committed Mar 20, 2024
1 parent 7f81fae commit 39a3cff
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 8 deletions.
15 changes: 11 additions & 4 deletions app.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import base64
import re
import io
import magic
from PIL import Image
from dotenv import load_dotenv
from quart import Quart, request, jsonify
Expand Down Expand Up @@ -170,24 +171,30 @@ async def image_validation():

# extension check
if image_extension not in valide_extension:
raise ImageValidationError("Invalid file extension")
raise ImageValidationError(f"invalid file extension: {image_extension}")

expected_header = f"data:image/{image_extension};base64"
expected_magic_header =f"image/{image_extension}"

# magic header check
magic_header = magic.from_buffer(base64.b64decode(encoded_image), mime=True)
if magic_header != expected_magic_header:
raise ImageValidationError(f"invalid file header: {magic_header}")

# header check
if header.lower() != expected_header:
raise ImageValidationError("Invalid file header")
raise ImageValidationError(f"invalid file header: {header}")

# size check
if image.size[0] > valide_dimension[0] and image.size[1] > valide_dimension[1]:
raise ImageValidationError("Invalid file size")
raise ImageValidationError(f"invalid file size: {image.size[0]}x{image.size[1]}")

# resizable check
try:
size = (100,150)
image.thumbnail(size)
except IOError:
raise ImageValidationError("Invalid file not resizable")
raise ImageValidationError("invalid file not resizable")

validator = await azure_storage_api.generate_hash(base64.b64decode(encoded_image))
CACHE['validators'].append(validator)
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@ quart-cors
python-dotenv
hypercorn
Pillow
python-magic
8 changes: 4 additions & 4 deletions tests/test_image_validation.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def test_invalid_header_image_validation(self):
data = json.loads(asyncio.run(response.get_data()))

self.assertEqual(response.status_code, 400)
self.assertEqual(data[0], 'Invalid file header')
self.assertEqual(data[0], 'invalid file header: data:image/')

@patch("PIL.Image.open")
def test_invalid_extension(self, mock_open):
Expand All @@ -78,7 +78,7 @@ def test_invalid_extension(self, mock_open):
data = json.loads(asyncio.run(response.get_data()))

self.assertEqual(response.status_code, 400)
self.assertEqual(data[0], 'Invalid file extension')
self.assertEqual(data[0], 'invalid file extension: md')

@patch("PIL.Image.open")
def test_invalid_size(self, mock_open):
Expand All @@ -103,7 +103,7 @@ def test_invalid_size(self, mock_open):
data = json.loads(asyncio.run(response.get_data()))

self.assertEqual(response.status_code, 400)
self.assertEqual(data[0], 'Invalid file size')
self.assertEqual(data[0], 'invalid file size: 2000x2000')

@patch("PIL.Image.open")
def test_rezisable_error(self, mock_open):
Expand All @@ -129,7 +129,7 @@ def test_rezisable_error(self, mock_open):
data = json.loads(asyncio.run(response.get_data()))

self.assertEqual(response.status_code, 400)
self.assertEqual(data[0], 'Invalid file not resizable')
self.assertEqual(data[0], 'invalid file not resizable')

if __name__ == '__main__':
unittest.main()

0 comments on commit 39a3cff

Please sign in to comment.