server/func: add content-type related utilities

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

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