From 1e754f229a7c4e495901659111950b886606c90f Mon Sep 17 00:00:00 2001 From: Yuxin Wu Date: Tue, 28 Jan 2025 05:54:01 -0800 Subject: [PATCH] add wxgf client code --- dump-html.py | 7 ++++-- requirements.txt | 1 + wechat/res.py | 57 ++++++++++++++++++++++++++++++++------------- wechat/wxgf.py | 60 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 107 insertions(+), 18 deletions(-) create mode 100644 wechat/wxgf.py diff --git a/dump-html.py b/dump-html.py index dbdea83..bb02bfd 100755 --- a/dump-html.py +++ b/dump-html.py @@ -17,8 +17,9 @@ def get_args(): parser.add_argument('name', help='name of contact') parser.add_argument('--output', help='output html file', default='output.html') parser.add_argument('--db', default='decoded.db', help='path to decoded database') - parser.add_argument('--avt', default='avatar.index', help='path to avatar.index file that only exists in old version of wechat') parser.add_argument('--res', default='resource', help='reseource directory') + parser.add_argument('--wxgf-server', help='address of the wxgf image decoder server') + parser.add_argument('--avt', default='avatar.index', help='path to avatar.index file that only exists in old version of wechat. Ignore for new version of wechat.') args = parser.parse_args() return args @@ -38,7 +39,9 @@ if __name__ == '__main__': sys.stderr.write(u"Couldn't find the chat {}.".format(name)); sys.exit(1) - res = Resource(parser, args.res, args.avt) + res = Resource(parser, args.res, + wxgf_server=args.wxgf_server, + avt_db=args.avt) msgs = parser.msgs_by_chat[chatid] logger.info(f"Number of Messages for chatid {chatid}: {len(msgs)}") assert len(msgs) > 0 diff --git a/requirements.txt b/requirements.txt index 4d7d41f..9771a6f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ pysqlcipher3>=1.0.3 csscompressor numpy ipython +websocket-client diff --git a/wechat/res.py b/wechat/res.py index 0e0cc33..b2d7bb2 100644 --- a/wechat/res.py +++ b/wechat/res.py @@ -1,9 +1,8 @@ # -*- coding: UTF-8 -*- -import glob import os -import re from PIL import Image +import time import io import base64 import logging @@ -15,9 +14,9 @@ import atexit from .emoji import EmojiReader from .avatar import AvatarReader from .common.textutil import md5 as get_md5_hex, get_file_b64 -from .common.timer import timing from .msg import TYPE_SPEAK from .audio import parse_wechat_audio_file +from .wxgf import WxgfAndroidDecoder, is_wxgf_file LIB_PATH = os.path.dirname(os.path.abspath(__file__)) VOICE_DIRNAME = 'voice2' @@ -28,11 +27,16 @@ VIDEO_DIRNAME = 'video' JPEG_QUALITY = 50 class Resource(object): - """ multimedia resources in chat""" - def __init__(self, parser, res_dir, avt_db): + """ Multimedia resources parser.""" + def __init__(self, parser, + res_dir: str, + *, + wxgf_server: str | None = None, + avt_db: str | None = None): """ Args: res_dir: path to the resource directory + wxgf_server: "hostname:port" that points to the wxgf converter android app avt_db: "avatar.index" file that only exists in old versions of wechat """ def check(subdir): @@ -47,6 +51,7 @@ class Resource(object): self.voice_dir = os.path.join(res_dir, VOICE_DIRNAME) 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) def _get_voice_filename(self, imgpath): @@ -146,20 +151,40 @@ class Resource(object): def get_jpg_b64(img_file): if not img_file: return None - if not img_file.endswith('jpg') and \ - imghdr.what(img_file) != 'jpeg': + + # True jpeg. Simplest case. + if img_file.endswith('jpg') and \ + imghdr.what(img_file) == 'jpeg': + return get_file_b64(img_file) + + if is_wxgf_file(img_file): + start = time.time() + buf = self.wxgf_decoder.decode_with_cache(img_file, None) + if buf is None: + if not self.wxgf_decoder.has_server(): + logger.warning("wxgf decoder server is not provided. Cannot decode wxgf images. Please follow instructions to create wxgf decoder server if these images need to be decoded.") + else: + logger.error("Failed to decode wxgf file: {}".format(img_file)) + return None + else: + elapsed = time.time() - start + if elapsed > 0.01 and self.wxgf_decoder.has_server(): + logger.info(f"Decoded {img_file} in {elapsed:.2f} seconds") + else: + with open(img_file, 'rb') as f: + buf = f.read() + + # File is not actually jpeg. Convert. + if imghdr.what(file=None, h=buf) != 'jpeg': try: - im = Image.open(open(img_file, 'rb')) + im = Image.open(io.BytesIO(buf)) except: return None - buf = io.BytesIO() - im.convert('RGB').save(buf, 'JPEG', quality=JPEG_QUALITY) - return base64.b64encode(buf.getvalue()).decode('ascii') - with open(img_file, 'rb') as f: - if f.read(4) == b'wxgf': - logger.warning(f"Don't know how to decode wxgf image {img_file}") - return None - return get_file_b64(img_file) + else: + bufio = io.BytesIO() + im.convert('RGB').save(bufio, 'JPEG', quality=JPEG_QUALITY) + buf = bufio.getvalue() + return base64.b64encode(buf).decode('ascii') big_file = get_jpg_b64(big_file) if big_file: diff --git a/wechat/wxgf.py b/wechat/wxgf.py new file mode 100644 index 0000000..980d5ec --- /dev/null +++ b/wechat/wxgf.py @@ -0,0 +1,60 @@ +from websocket import create_connection +import os + + +WXGF_HEADER = b'wxgf' +FAILURE_MESSAGE = b'FAILED' + + +class WxgfAndroidDecoder: + + def __init__(self, server: str | None): + """server: hostname:port""" + if server is not None: + if "://" not in server: + server = "ws://" + server + self.ws = create_connection(server) + + def __del__(self): + self.ws.close() + + def has_server(self) -> bool: + return hasattr(self, 'ws') + + def decode(self, data: bytes) -> bytes | None: + assert data[:4] == WXGF_HEADER, data[:20] + self.ws.send(data, opcode=0x2) + res = self.ws.recv() + if res == FAILURE_MESSAGE: + return None + return res + + def decode_with_cache(self, fname: str, data: bytes | None) -> bytes | None: + """Decode and save cache. + + Args: + fname: original file path. cache will be saved alongside. + data: data to decode. None to use content of fname. + """ + if data is None: + with open(fname, 'rb') as f: + data = f.read() + out_fname = os.path.splitext(fname)[0] + '.dec' + + if os.path.exists(out_fname): + with open(out_fname, 'rb') as f: + return f.read() + + if not self.has_server(): + return None + res = self.decode(data) + + if res is not None: + with open(out_fname, 'wb') as f: + f.write(res) + return res + + +def is_wxgf_file(fname): + with open(fname, 'rb') as f: + return f.read(4) == WXGF_HEADER