diff --git a/dump-html.py b/dump-html.py index f7f16d0..1a2e35b 100755 --- a/dump-html.py +++ b/dump-html.py @@ -3,6 +3,7 @@ import os import sys import argparse import logging +from datetime import datetime from wechat.parser import WeChatDBParser from wechat.res import Resource @@ -19,6 +20,8 @@ def get_args(): parser.add_argument('--res', default='resource', help='the resource 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.') + parser.add_argument('--start', help='start time in format of YYYY-MM-DD HH:MM:SS', + type=datetime.fromisoformat) args = parser.parse_args() return args @@ -43,6 +46,9 @@ if __name__ == '__main__': msgs = parser.msgs_by_chat[chatid] logger.info(f"Number of Messages for chatid {chatid}: {len(msgs)}") assert len(msgs) > 0 + if args.start is not None: + msgs = [msg for msg in msgs if msg.createTime > args.start] + logger.info(f"Number of Messages after {args.start}: {len(msgs)}") render = HTMLRender(parser, res) htmls = render.render_msgs(msgs) diff --git a/wechat/avatar.py b/wechat/avatar.py index 2c9089a..6099f65 100644 --- a/wechat/avatar.py +++ b/wechat/avatar.py @@ -49,7 +49,7 @@ class AvatarReader(object): except Exception: pass - def get_avatar_from_avtdir(self, avtid): + def get_avatar_from_avtdir(self, avtid) -> Image.Image | None: dir1, dir2 = avtid[:2], avtid[2:4] candidates = glob.glob(os.path.join(self.avt_dir, dir1, dir2, f"*{avtid}*")) candidates = sorted(set(candidates), key=_filename_priority, reverse=True) @@ -83,7 +83,16 @@ class AvatarReader(object): ret = self.get_avatar_from_avtdir(avtid) if ret is not None: return ret - logger.warning("Avatar for {} not found anywhere.".format(username)) + logger.warning("Avatar file for {} not found.".format(username)) + + def save_avatar_to_avtdir(self, username: str, im: Image.Image): + """Save a downloaded avatar to avtdir so it can be reused next time.""" + avtid = md5(username.encode('utf-8')) + dir1, dir2 = avtid[:2], avtid[2:4] + fname = os.path.join(self.avt_dir, dir1, dir2, f"user_{avtid}.png") + os.makedirs(os.path.dirname(fname), exist_ok=True) + logger.info(f"Caching downloaded avatar for {username} to {fname}.") + im.save(fname, 'PNG') def read_img_from_block(self, filename, pos, size): file_idx = pos >> 32 diff --git a/wechat/msg.py b/wechat/msg.py index cf0df83..56614d2 100644 --- a/wechat/msg.py +++ b/wechat/msg.py @@ -155,6 +155,12 @@ class WeChatMsg(object): else: return ret + def __eq__(self, r): + return self.createTime == r.createTime and \ + self.talker == r.talker and \ + self.isSend == r.isSend + # imgPath might change after migration. + def __lt__(self, r): return self.createTime < r.createTime diff --git a/wechat/parser.py b/wechat/parser.py index 662848e..1fa17a0 100644 --- a/wechat/parser.py +++ b/wechat/parser.py @@ -16,6 +16,7 @@ addr_upload2 chatroom message rcontact +img_flag """ class WeChatDBParser(object): @@ -36,6 +37,7 @@ class WeChatDBParser(object): self.emoji_groups = {} self.emoji_info = {} self.emoji_encryption_key = None + self.avatar_urls = {} self._parse() def _parse_contact(self): @@ -126,6 +128,14 @@ SELECT {} FROM message if cdnUrl or encrypturl: self.emoji_info[md5] = (catalog, cdnUrl, encrypturl, aeskey) + def _parse_img_flag(self): + """Parse the img_flag table which stores avatar for each id.""" + query = self.cc.execute( +""" SELECT username, reserved1 FROM img_flag """) + for row in query: + username, url = row + if url: + self.avatar_urls[username] = url def _parse(self): self._parse_contact() @@ -133,6 +143,7 @@ SELECT {} FROM message self._parse_msg() self._parse_imginfo() self._parse_emoji() + self._parse_img_flag() def get_emoji_encryption_key(self): # obtain local encryption key in a special entry in the database diff --git a/wechat/res.py b/wechat/res.py index e3810bc..d74ba78 100644 --- a/wechat/res.py +++ b/wechat/res.py @@ -2,6 +2,7 @@ import os from PIL import Image +import requests import time import io import base64 @@ -83,12 +84,24 @@ class Resource(object): self.voice_cache = [pool.apply_async(parse_wechat_audio_file, (self._get_voice_filename(k),)) for k in voice_paths] - def get_avatar(self, username): + def get_avatar(self, username) -> str: """ return base64 unicode string""" im = self.avt_reader.get_avatar(username) if im is None: - logger.warning(f"Cannot find avatar for {username}.") - return "" + # Try downloading the avatar directly. + avatar_url = self.parser.avatar_urls.get(username) + if avatar_url is None: + return "" + logger.info(f"Requesting avatar of {username} from {avatar_url} ...") + try: + r = requests.get(avatar_url).content + im = Image.open(io.BytesIO(r)) + except Exception: + logger.exception(f"Failed to fetch avatar of {username}.") + return "" + else: + self.avt_reader.save_avatar_to_avtdir(username, im) + buf = io.BytesIO() try: im.save(buf, 'JPEG', quality=JPEG_QUALITY)