read avatar from sfs/avatar.block.0000x

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,7 +28,7 @@ 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