add wxgf client code

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

requirements.txt

@@ -7,3 +7,4 @@ pysqlcipher3>=1.0.3
 csscompressor
 numpy
 ipython
+websocket-client

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:

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