diff --git a/wechat/common/imgutil.py b/wechat/common/imgutil.py new file mode 100644 index 0000000..167a547 --- /dev/null +++ b/wechat/common/imgutil.py @@ -0,0 +1,51 @@ +import io +from pathlib import Path + +from PIL import Image + +try: + import imghdr as _imghdr +except ModuleNotFoundError: # Python 3.13+ + _imghdr = None + + +def what(file: str | Path | None = None, h: bytes | None = None) -> str | None: + """ + Compatibility wrapper for the removed stdlib `imghdr` module (Python 3.13+). + + Behaves like `imghdr.what(file, h)` and returns a lowercase format string + like "jpeg", "png", "gif", or None if unknown. + """ + if _imghdr is not None: + return _imghdr.what(file, h) + if h is not None: + return _what_from_bytes(h) + if file is None: + return None + try: + with Image.open(str(file)) as im: + fmt = im.format + except Exception: + return None + return fmt.lower() if fmt else None + + +def _what_from_bytes(data: bytes) -> str | None: + if data.startswith(b"\xff\xd8\xff"): + return "jpeg" + if data.startswith(b"\x89PNG\r\n\x1a\n"): + return "png" + if data[:6] in (b"GIF87a", b"GIF89a"): + return "gif" + if data.startswith(b"BM"): + return "bmp" + if len(data) >= 12 and data.startswith(b"RIFF") and data[8:12] == b"WEBP": + return "webp" + if data.startswith(b"II*\x00") or data.startswith(b"MM\x00*"): + return "tiff" + try: + with Image.open(io.BytesIO(data)) as im: + fmt = im.format + except Exception: + return None + return fmt.lower() if fmt else None diff --git a/wechat/emoji.py b/wechat/emoji.py index 8819e82..3964dbd 100644 --- a/wechat/emoji.py +++ b/wechat/emoji.py @@ -4,13 +4,13 @@ import logging import io import requests import base64 -import imghdr from PIL import Image import pickle from Crypto.Cipher import AES from .wxgf import WxgfAndroidDecoder, is_wxgf_buffer from .parser import WeChatDBParser +from .common.imgutil import what as img_what from .common.textutil import md5 as get_md5_hex, get_file_b64, get_file_md5 @@ -124,10 +124,10 @@ class EmojiReader: candidates = [] def get_data_no_fallback(fname): - if imghdr.what(fname): + if img_what(fname): data_md5 = get_file_md5(fname) if data_md5 == md5: - return get_file_b64(fname), imghdr.what(fname) + return get_file_b64(fname), img_what(fname) try: content = self._decode_emoji(fname) @@ -147,9 +147,9 @@ class EmojiReader: logger.error(f"Error decoding emoji {fname} : {str(e)}") def get_data_fallback(fname): - if not imghdr.what(fname): + if not img_what(fname): return # fallback files are not encrypted - return get_file_b64(fname), imghdr.what(fname) + return get_file_b64(fname), img_what(fname) get_data_func = get_data_fallback if allow_fallback else get_data_no_fallback results = [(x, get_data_func(x)) for x in candidates] diff --git a/wechat/res.py b/wechat/res.py index d74ba78..70d08b7 100644 --- a/wechat/res.py +++ b/wechat/res.py @@ -8,12 +8,12 @@ import io import base64 import logging logger = logging.getLogger(__name__) -import imghdr from multiprocessing import Pool import atexit from .emoji import EmojiReader from .avatar import AvatarReader +from .common.imgutil import what as img_what from .common.textutil import md5 as get_md5_hex, get_file_b64 from .msg import TYPE_SPEAK from .audio import parse_wechat_audio_file @@ -167,7 +167,7 @@ class Resource(object): # True jpeg. Simplest case. if img_file.endswith('jpg') and \ - imghdr.what(img_file) == 'jpeg': + img_what(img_file) == 'jpeg': return get_file_b64(img_file) if is_wxgf_file(img_file): @@ -188,7 +188,7 @@ class Resource(object): buf = f.read() # File is not actually jpeg. Convert. - if imghdr.what(file=None, h=buf) != 'jpeg': + if img_what(file=None, h=buf) != 'jpeg': try: im = Image.open(io.BytesIO(buf)) except: