diff --git a/server/szurubooru/func/mime.py b/server/szurubooru/func/mime.py new file mode 100644 index 00000000..a7d26ca9 --- /dev/null +++ b/server/szurubooru/func/mime.py @@ -0,0 +1,49 @@ +import re + +def get_mime_type(content): + if not content: + return None + + if content[0:3] in (b'CWS', b'FWS', b'ZWS'): + return 'application/x-shockwave-flash' + + if content[0:3] == b'\xFF\xD8\xFF': + return 'image/jpeg' + + if content[0:6] == b'\x89PNG\x0D\x0A': + return 'image/png' + + if content[0:6] in (b'GIF87a', b'GIF89a'): + return 'image/gif' + + if content[0:4] == b'\x1A\x45\xDF\xA3': + return 'video/webm' + + if content[4:12] in (b'ftypisom', b'ftypmp42'): + return 'video/mp4' + + return 'application/octet-stream' + +def get_extension(mime_type): + extension_map = { + 'application/x-shockwave-flash': 'swf', + 'image/gif': 'gif', + 'image/jpeg': 'jpg', + 'image/png': 'png', + 'video/mp4': 'mp4', + 'video/webm': 'webm', + } + return extension_map.get(mime_type.strip().lower(), None) + +def is_flash(mime_type): + return mime_type.lower() == 'application/x-shockwave-flash' + +def is_video(mime_type): + return mime_type.lower() in ('application/ogg', 'video/mp4', 'video/webm') + +def is_image(mime_type): + return mime_type.lower() in ('image/jpeg', 'image/png', 'image/gif') + +def is_animated_gif(content): + return get_mime_type(content) == 'image/gif' \ + and len(re.findall(b'\x21\xF9\x04.{4}\x00[\x2C\x21]', content)) > 1 diff --git a/server/szurubooru/tests/assets/flash.swf b/server/szurubooru/tests/assets/flash.swf new file mode 100644 index 00000000..c6195c41 Binary files /dev/null and b/server/szurubooru/tests/assets/flash.swf differ diff --git a/server/szurubooru/tests/assets/gif-animated.gif b/server/szurubooru/tests/assets/gif-animated.gif new file mode 100644 index 00000000..ce8c95b9 Binary files /dev/null and b/server/szurubooru/tests/assets/gif-animated.gif differ diff --git a/server/szurubooru/tests/assets/gif.gif b/server/szurubooru/tests/assets/gif.gif new file mode 100644 index 00000000..edaf2b97 Binary files /dev/null and b/server/szurubooru/tests/assets/gif.gif differ diff --git a/server/szurubooru/tests/assets/jpeg.jpg b/server/szurubooru/tests/assets/jpeg.jpg new file mode 100644 index 00000000..71911bf4 Binary files /dev/null and b/server/szurubooru/tests/assets/jpeg.jpg differ diff --git a/server/szurubooru/tests/assets/mp4.mp4 b/server/szurubooru/tests/assets/mp4.mp4 new file mode 100644 index 00000000..3c76f5c9 Binary files /dev/null and b/server/szurubooru/tests/assets/mp4.mp4 differ diff --git a/server/szurubooru/tests/assets/png-transparent.png b/server/szurubooru/tests/assets/png-transparent.png new file mode 100644 index 00000000..91a99b94 Binary files /dev/null and b/server/szurubooru/tests/assets/png-transparent.png differ diff --git a/server/szurubooru/tests/assets/webm.webm b/server/szurubooru/tests/assets/webm.webm new file mode 100644 index 00000000..38e17b83 Binary files /dev/null and b/server/szurubooru/tests/assets/webm.webm differ diff --git a/server/szurubooru/tests/func/test_mime.py b/server/szurubooru/tests/func/test_mime.py new file mode 100644 index 00000000..830b275a --- /dev/null +++ b/server/szurubooru/tests/func/test_mime.py @@ -0,0 +1,54 @@ +import os +import pytest +from szurubooru.func import mime + +@pytest.mark.parametrize('input_path,expected_mime_type', [ + ('mp4.mp4', 'video/mp4'), + ('webm.webm', 'video/webm'), + ('flash.swf', 'application/x-shockwave-flash'), + ('png-transparent.png', 'image/png'), + ('jpeg.jpg', 'image/jpeg'), + ('gif.gif', 'image/gif'), +]) +def test_get_mime_type(read_asset, input_path, expected_mime_type): + assert mime.get_mime_type(read_asset(input_path)) == expected_mime_type + +@pytest.mark.parametrize('input_mime_type,expected_state', [ + ('application/x-shockwave-flash', True), + ('APPLICATION/X-SHOCKWAVE-FLASH', True), + ('application/x-shockwave', False), +]) +def test_is_flash(input_mime_type, expected_state): + assert mime.is_flash(input_mime_type) == expected_state + +@pytest.mark.parametrize('input_mime_type,expected_state', [ + ('video/webm', True), + ('VIDEO/WEBM', True), + ('video/mp4', True), + ('VIDEO/MP4', True), + ('video/anything_else', False), + ('application/ogg', True), + ('not a video', False), +]) +def test_is_video(input_mime_type, expected_state): + assert mime.is_video(input_mime_type) == expected_state + +@pytest.mark.parametrize('input_mime_type,expected_state', [ + ('image/gif', True), + ('image/png', True), + ('image/jpeg', True), + ('IMAGE/GIF', True), + ('IMAGE/PNG', True), + ('IMAGE/JPEG', True), + ('image/anything_else', False), + ('not an image', False), +]) +def test_is_image(input_mime_type, expected_state): + assert mime.is_image(input_mime_type) == expected_state + +@pytest.mark.parametrize('input_path,expected_state', [ + ('gif.gif', False), + ('gif-animated.gif', True), +]) +def test_is_animated_gif(read_asset, input_path, expected_state): + assert mime.is_animated_gif(read_asset(input_path)) == expected_state