Merge pull request #14 from SwordYork/master

Avatar support for new version wechat

dump-html.py

@@ -5,8 +5,8 @@
 # Author: Yuxin Wu <[email protected]>
 
 import sys
-if len(sys.argv) != 5:
-    sys.exit("Usage: {0} <path to decoded_database.db> <path to resource> <name> <output html>".format(sys.argv[0]))
+if len(sys.argv) != 6:
+    sys.exit("Usage: {0} <path to decoded_database.db> <path to avatar.index> <path to resource> <name> <output html>".format(sys.argv[0]))
 
 from common.textutil import ensure_unicode
 from wechat.parser import WeChatDBParser
@@ -14,12 +14,13 @@ from wechat.res import Resource
 from wechat.render import HTMLRender
 
 db_file = sys.argv[1]
-resource_dir = sys.argv[2]
-name = ensure_unicode(sys.argv[3])
-output_file = sys.argv[4]
+avt_db = sys.argv[2]
+resource_dir = sys.argv[3]
+name = ensure_unicode(sys.argv[4])
+output_file = sys.argv[5]
 
 parser = WeChatDBParser(db_file)
-res = Resource(resource_dir)
+res = Resource(resource_dir, avt_db)
 
 try:
     msgs = parser.msgs_by_talker[name]

wechat/avatar.py

@@ -7,28 +7,59 @@
 import os
 import numpy as np
 import logging
+import sqlite3
 logger = logging.getLogger(__name__)
 
 from common.textutil import ensure_bin_str, md5
 
 class AvatarReader(object):
-    def __init__(self, avt_dir):
+    def __init__(self, avt_dir, avt_db="avatar.index"):
         self.avt_dir = avt_dir
+        self.avt_db = avt_db
 
     def get_avatar(self, username):
         """ username: `username` field in db.rcontact"""
         username = ensure_bin_str(username)
         filename = md5(username)
         dir1, dir2 = filename[:2], filename[2:4]
-        filename = os.path.join(self.avt_dir, dir1, dir2,
+        filename = os.path.join(dir1, dir2,
                                 "user_{}.png.bm".format(filename))
-        if not os.path.isfile(filename):
+
+        index_avatar = self.query_index(filename)
+        if index_avatar == -1:
             logger.warn("Avatar not found for {}".format(username))
             return None
         else:
-            img = AvatarReader.read_bm(filename)
+            img = self.read_bm_block(index_avatar)
             return img
 
+    def read_bm_block(self, pos):
+        hex_pos = hex(pos)
+        fname = os.path.join(self.avt_dir, 
+                'avatar.block.0000' + hex_pos[2])
+        f = open(fname, 'rb')
+        start_pos = pos - 2 ** 34
+        f.seek(start_pos+30)
+        while f.read(1) != b"\x00":
+            continue
+
+        size = (96, 96, 3)
+        img = np.zeros(size, dtype='uint8')
+        for i in range(96):
+            for j in range(96):
+                r, g, b, a = map(ord, f.read(4))
+                img[i,j] = (r, g, b)
+        return img
+
+    def query_index(self, filename):
+        conn = sqlite3.connect(self.avt_db)
+        try:
+            cursor = conn.execute("select Offset from Index_avatar where FileName='{}'".format(filename))
+            index_avatar = cursor.fetchone()[0]
+            return index_avatar
+        except:
+            return -1
+
     @staticmethod
     def read_bm(fname):
         size = (96, 96, 3)

wechat/res.py

@@ -28,13 +28,13 @@ INTERNAL_EMOJI_DIR = os.path.join(LIB_PATH, 'static', 'internal_emoji')
 VOICE_DIRNAME = 'voice2'
 IMG_DIRNAME = 'image2'
 EMOJI_DIRNAME = 'emoji'
-AVATAR_DIRNAME = 'avatar'
+AVATAR_DIRNAME = 'sfs'
 
 JPEG_QUALITY = 50
 
 class Resource(object):
     """ multimedia resources in chat"""
-    def __init__(self, res_dir):
+    def __init__(self, res_dir, avt_db):
         def check(subdir):
             assert os.path.isdir(os.path.join(res_dir, subdir)), \
                     "No such directory: {}".format(subdir)
@@ -45,7 +45,7 @@ class Resource(object):
         self.img_dir = os.path.join(res_dir, IMG_DIRNAME)
         self.voice_dir = os.path.join(res_dir, VOICE_DIRNAME)
         self.emoji_dir = os.path.join(res_dir, EMOJI_DIRNAME)
-        self.avt_reader = AvatarReader(os.path.join(res_dir, AVATAR_DIRNAME))
+        self.avt_reader = AvatarReader(os.path.join(res_dir, AVATAR_DIRNAME), avt_db)
 
     def get_voice_filename(self, imgpath):
         fname = md5(imgpath)
@@ -140,7 +140,7 @@ class Resource(object):
                imghdr.what(img_file) != 'jpeg':
                 im = Image.open(open(img_file, 'rb'))
                 buf = cStringIO.StringIO()
-                im.save(buf, 'JPEG', quality=JPEG_QUALITY)
+                im.convert('RGB').save(buf, 'JPEG', quality=JPEG_QUALITY)
                 return base64.b64encode(buf.getvalue())
             return get_file_b64(img_file)
         big_file = get_jpg_b64(big_file)