From b271469b94e35f2c584ab34156c504d377744970 Mon Sep 17 00:00:00 2001 From: Yuxin Wu Date: Wed, 29 Jan 2025 15:07:52 -0800 Subject: [PATCH] support wxgf emoji --- dump-html.py | 1 - wechat/emoji.py | 24 ++++++++++----- wechat/render.py | 45 +++++++++++++++++++---------- wechat/res.py | 7 ++--- wechat/static/TP_QQMUSIC_NOIMG.html | 15 ++++++++++ wechat/wxgf.py | 3 ++ 6 files changed, 67 insertions(+), 28 deletions(-) create mode 100644 wechat/static/TP_QQMUSIC_NOIMG.html diff --git a/dump-html.py b/dump-html.py index 7cc8a65..0219d65 100755 --- a/dump-html.py +++ b/dump-html.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -# -*- coding: UTF-8 -*- import os import sys import argparse diff --git a/wechat/emoji.py b/wechat/emoji.py index bf9be5f..8819e82 100644 --- a/wechat/emoji.py +++ b/wechat/emoji.py @@ -1,7 +1,6 @@ import os from pathlib import Path import logging -import tempfile import io import requests import base64 @@ -10,6 +9,7 @@ from PIL import Image import pickle from Crypto.Cipher import AES +from .wxgf import WxgfAndroidDecoder, is_wxgf_buffer from .parser import WeChatDBParser from .common.textutil import md5 as get_md5_hex, get_file_b64, get_file_md5 @@ -23,21 +23,20 @@ def _get_aes_key(md5): # ascii representation of the first half of md5 is used as aes key assert len(md5) == 32 return md5[:16].encode('ascii') - # ret = "" - # for ch in md5[:16]: - # ret += format(ord(ch), 'x') - # return ret class EmojiReader: def __init__(self, resource_dir: str, parser: WeChatDBParser, + *, + wxgf_decoder: WxgfAndroidDecoder, cache_file: str=None): """ Args: resource_dir: path to resource/ parser: Database parser + wxgf_decoder: Wxgf image decoder cache_file: a cache file to store emoji downloaded from URLs. default to a emoji.cache file under wechat-dump. """ @@ -48,6 +47,7 @@ class EmojiReader: # mapping from md5 to the (cdnurl, encrypturl, aeskey) # columns in EmojiInfo table. self.cache_file = cache_file or DEFAULT_EMOJI_CACHE + self.wxgf_decoder = wxgf_decoder # cache stores md5 -> (base64str, format) if os.path.isfile(self.cache_file): @@ -133,9 +133,14 @@ class EmojiReader: content = self._decode_emoji(fname) data_md5 = get_md5_hex(content) if data_md5 != md5: - if content.startswith(b"wxgf"): - raise ValueError("Unsupported mysterious image format: wxgf") - raise ValueError("Decoded data mismatch md5!") + if is_wxgf_buffer(content): + content = self.wxgf_decoder.decode_with_cache(fname, content) + if content is None: + if not self.wxgf_decoder.has_server(): + logger.warning("wxgf decoder server is not provided. Cannot decode wxgf emojis.") + raise ValueError("Failed to decode wxgf file.") + else: + raise ValueError("Decoded data mismatch md5!") im = Image.open(io.BytesIO(content)) return (base64.b64encode(content).decode('ascii'), im.format.lower()) except Exception as e: @@ -183,6 +188,9 @@ class EmojiReader: try: logger.info("Requesting encrypted emoji {} from {} ...".format(md5, encrypturl)) buf = requests.get(encrypturl).content + if buf == b'': + logger.error(f"Failed to download emoji {md5}") + return None, None aeskey = bytes.fromhex(aeskey) cipher = AES.new(aeskey, AES.MODE_CBC, iv=aeskey) decoded_buf = cipher.decrypt(buf) diff --git a/wechat/render.py b/wechat/render.py index 4c78f66..bfcc0f1 100644 --- a/wechat/render.py +++ b/wechat/render.py @@ -3,6 +3,7 @@ import os from collections import Counter +from functools import lru_cache import glob from pyquery import PyQuery import logging @@ -35,10 +36,19 @@ TEMPLATES_FILES = {TYPE_MSG: "TP_MSG", TYPE_VIDEO_FILE: "TP_VIDEO_FILE", TYPE_QQMUSIC: "TP_QQMUSIC", } -TEMPLATES = { - k: open(os.path.join(STATIC_PATH, '{}.html'.format(v))).read() - for k, v in TEMPLATES_FILES.items() -} + + +@lru_cache() +def get_template(name: str | int) -> str | None: + """Return the html template given a file name or msg type.""" + if isinstance(name, int): + name = TEMPLATES_FILES.get(name, None) + if name is None: + return None + html_path = os.path.join(STATIC_PATH, f"{name}.html") + with open(html_path) as f: + return f.read() + class HTMLRender(object): def __init__(self, parser, res=None): @@ -107,7 +117,7 @@ class HTMLRender(object): format_dict['nickname'] = ' ' def fallback(): - template = TEMPLATES[TYPE_MSG] + template = get_template(TYPE_MSG) content = msg.msg_str() content = self.smiley.replace_smileycode(content) if not msg.known_type: @@ -115,7 +125,7 @@ class HTMLRender(object): content = html.escape(content) return template.format(content=content, **format_dict) - template = TEMPLATES.get(msg.type) + template = get_template(msg.type) if msg.type == TYPE_SPEAK: audio_str, duration = self.res.get_voice_mp3(msg.imgPath) format_dict['voice_duration'] = duration @@ -140,10 +150,13 @@ class HTMLRender(object): jobj = json.loads(msg.msg_str()) content = f"{jobj['title']} - {jobj['singer']}" - # imgPath was original THUMBNAIL_DIRPATH://th_xxxxxxxxx - imgpath = msg.imgPath.split('_')[-1] - img = self.res.get_img([imgpath]) - format_dict['img'] = (img, 'jpeg') + if msg.imgPath is not None: + # imgPath was original THUMBNAIL_DIRPATH://th_xxxxxxxxx + imgpath = msg.imgPath.split('_')[-1] + img = self.res.get_img([imgpath]) + format_dict['img'] = (img, 'jpeg') + else: + template = get_template("TP_QQMUSIC_NOIMG") return template.format(url=jobj['url'], content=content, **format_dict) elif msg.type == TYPE_EMOJI or msg.type == TYPE_CUSTOM_EMOJI: if 'emoticonmd5' in msg.content: @@ -171,7 +184,12 @@ class HTMLRender(object): return template.format(**format_dict) elif msg.type == TYPE_VIDEO_FILE: video = self.res.get_video(msg.imgPath) - if video.endswith(".mp4"): + if video is None: + logger.warning(f"Cannot find video {msg.imgPath} ({msg.createTime})") + # fallback + format_dict['content'] = f"VIDEO FILE {msg.imgPath}" + return get_template(TYPE_MSG).format(**format_dict) + elif video.endswith(".mp4"): video_str = get_file_b64(video) format_dict["video_str"] = video_str return template.format(**format_dict) @@ -179,10 +197,7 @@ class HTMLRender(object): # only has thumbnail image_str = get_file_b64(video) format_dict["img"] = (image_str, 'jpeg') - return TEMPLATES[TYPE_IMG].format(**format_dict) - # fallback - format_dict['content'] = f"VIDEO FILE {msg.imgPath}" - return TEMPLATES_FILES[TYPE_MSG].format(**format_dict) + return get_template(TYPE_IMG).format(**format_dict) elif msg.type == TYPE_WX_VIDEO: # TODO: fetch video from resource return fallback() diff --git a/wechat/res.py b/wechat/res.py index b2d7bb2..e3810bc 100644 --- a/wechat/res.py +++ b/wechat/res.py @@ -52,7 +52,7 @@ class Resource(object): self.video_dir = os.path.join(res_dir, VIDEO_DIRNAME) self.avt_reader = AvatarReader(res_dir, avt_db) self.wxgf_decoder = WxgfAndroidDecoder(wxgf_server) - self.emoji_reader = EmojiReader(res_dir, self.parser) + self.emoji_reader = EmojiReader(res_dir, self.parser, wxgf_decoder=self.wxgf_decoder) def _get_voice_filename(self, imgpath): fname = get_md5_hex(imgpath.encode('ascii')) @@ -195,12 +195,11 @@ class Resource(object): """ Returns: (b64 encoded img string, format) """ return self.emoji_reader.get_emoji(md5) - def get_video(self, videoid): + def get_video(self, videoid) -> str | None: video_file = os.path.join(self.video_dir, videoid + ".mp4") video_thumbnail_file = os.path.join(self.video_dir, videoid + ".jpg") if os.path.exists(video_file): return video_file elif os.path.exists(video_thumbnail_file): return video_thumbnail_file - logger.warning(f"Cannot find video {videoid}") - return "" + return None diff --git a/wechat/static/TP_QQMUSIC_NOIMG.html b/wechat/static/TP_QQMUSIC_NOIMG.html new file mode 100644 index 0000000..26234e2 --- /dev/null +++ b/wechat/static/TP_QQMUSIC_NOIMG.html @@ -0,0 +1,15 @@ +
+
+ +
+
+
+
+ 
+QQ Music:{content}
+
+
+
+
+
+
diff --git a/wechat/wxgf.py b/wechat/wxgf.py index d3f1063..8b4e581 100644 --- a/wechat/wxgf.py +++ b/wechat/wxgf.py @@ -59,3 +59,6 @@ class WxgfAndroidDecoder: def is_wxgf_file(fname): with open(fname, 'rb') as f: return f.read(4) == WXGF_HEADER + +def is_wxgf_buffer(buf: bytes): + return buf[:4] == WXGF_HEADER