diff --git a/wechat/avatar.py b/wechat/avatar.py index 68776b0..e8ea5f2 100644 --- a/wechat/avatar.py +++ b/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) diff --git a/wechat/res.py b/wechat/res.py index 72a9e4a..5b818cd 100644 --- a/wechat/res.py +++ b/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