diff --git a/emails/store/file.py b/emails/store/file.py index d601f16..b1f94c6 100644 --- a/emails/store/file.py +++ b/emails/store/file.py @@ -2,6 +2,7 @@ import uuid from mimetypes import guess_type +import puremagic from email.mime.base import MIMEBase from email.encoders import encode_base64 from os.path import basename @@ -102,6 +103,23 @@ def get_mime_type(self) -> str: filename = self.filename if filename: r = self._mime_type = guess_type(filename)[0] + if not r: + _data = self._data + if isinstance(_data, bytes): + header = _data + elif isinstance(_data, str): + header = _data.encode() + elif _data is not None: + pos = _data.tell() + header = _data.read(128) + _data.seek(pos) + else: + header = None + if header: + try: + r = puremagic.from_string(header, mime=True) + except puremagic.PureError: + pass if not r: r = MIMETYPE_UNKNOWN self._mime_type = r diff --git a/emails/testsuite/store/test_store.py b/emails/testsuite/store/test_store.py index c5596f2..bedbb45 100644 --- a/emails/testsuite/store/test_store.py +++ b/emails/testsuite/store/test_store.py @@ -68,6 +68,33 @@ def test_get_data_none(): assert f.data is None +def test_mime_type_from_content(): + # PNG magic bytes, no file extension + png_header = (b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR' + b'\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02' + b'\x00\x00\x00\x90wS\xde') + f = BaseFile(data=png_header, filename='image_no_ext') + assert f.mime_type == 'image/png' + + # JPEG magic bytes, no file extension + jpeg_header = b'\xff\xd8\xff\xe0\x00\x10JFIF' + f = BaseFile(data=jpeg_header, filename='photo') + assert f.mime_type == 'image/jpeg' + + # Unknown bytes, no extension — should fall back to unknown + f = BaseFile(data=b'\x00\x01\x02\x03', filename='mystery') + assert f.mime_type == 'application/unknown' + + # Extension still takes priority + f = BaseFile(data=png_header, filename='image.gif') + assert f.mime_type == 'image/gif' + + # File-like data: mime detected without exhausting stream + f = BaseFile(data=BytesIO(png_header), filename='no_ext') + assert f.mime_type == 'image/png' + assert f.data == png_header # stream not consumed by mime detection + + def test_store_commons2(): store = emails.store.MemoryFileStore() f1 = store.add({'uri': '/a/c.gif'}) diff --git a/requirements/base.txt b/requirements/base.txt index efd3de4..466b8f5 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -4,3 +4,4 @@ chardet python-dateutil requests premailer>=2.8.3 +puremagic diff --git a/setup.py b/setup.py index 6ffea53..35801c6 100644 --- a/setup.py +++ b/setup.py @@ -127,7 +127,7 @@ def find_version(*file_paths): package_data={'emails': ['py.typed']}, scripts=['scripts/make_rfc822.py'], python_requires='>=3.10', - install_requires=['python-dateutil'], + install_requires=['python-dateutil', 'puremagic'], extras_require={ 'html': ['cssutils', 'lxml', 'chardet', 'requests', 'premailer'], },