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