fix old avatar storage

wechat/avatar.py

@@ -13,6 +13,13 @@ logger = logging.getLogger(__name__)
 from .common.textutil import ensure_unicode, md5
 
 
+def _filename_priority(s):
+    if "_hd" in s and s.endswith(".png"):
+        return 10
+    else:
+        return 1
+
+
 class AvatarReader(object):
     def __init__(self, res_dir, avt_db="avatar.index"):
         self.sfs_dir = os.path.join(res_dir, 'sfs')
@@ -22,14 +29,15 @@ class AvatarReader(object):
         self.avt_db = avt_db
         self._use_avt = True
         if os.path.isdir(self.avt_dir) and len(os.listdir(self.avt_dir)):
-            self.avt_use_db = False
+            # has resource/avatar
+            pass
         elif self.avt_db is not None \
                 and os.path.isfile(self.avt_db) \
                 and glob.glob(os.path.join(self.sfs_dir, 'avatar*')):
-            self.avt_use_db = True
+            # has sfs/avatar*
+            pass
         else:
-            logger.warn(
-                    "Cannot find avatar files. Will not use avatar!")
+            logger.warn("Cannot find avatar storage. Will not use avatar!")
             self._use_avt = False
 
     def get_avatar(self, username):
@@ -38,41 +46,39 @@ class AvatarReader(object):
             return None
         username = ensure_unicode(username)
         avtid = md5(username.encode('utf-8'))
+        try:
+            candidates = self.query_index(avtid)
+            candidates = sorted(candidates, key=lambda x: _filename_priority(x[0]), reverse=True)
+            for c in candidates:
+                path, offset, size = c
+                return self.read_img_from_block(path, offset, size)
+        except:
+            pass
+
         dir1, dir2 = avtid[:2], avtid[2:4]
         candidates = glob.glob(os.path.join(self.avt_dir, dir1, dir2, f"*{avtid}*"))
         default_candidate = os.path.join(self.avt_dir, dir1, dir2, f"user_{avtid}.png")
         candidates.append(default_candidate)
 
-        def priority(s):
-            if "_hd" in s and s.endswith(".png"):
-                return 10
-            else:
-                return 1
-
-        candidates = sorted(set(candidates), key=priority, reverse=True)
-
+        candidates = sorted(set(candidates), key=_filename_priority, reverse=True)
         for cand in candidates:
             try:
-                if self.avt_use_db:
-                    pos, size = self.query_index(cand)
-                    return self.read_img(pos, size)
-                else:
-                    if os.path.exists(cand):
-                        if cand.endswith(".bm"):
-                            return self.read_bm_file(cand)
-                        else:
-                            return Image.open(cand)
+                if os.path.exists(cand):
+                    if cand.endswith(".bm"):
+                        return self.read_bm_file(cand)
+                    else:
+                        return Image.open(cand)
             except Exception:
-                logger.exception("HHH")
+                logger.exception("")
                 pass
-        logger.warning("Avatar for {} not found in avatar database.".format(username))
+        logger.warning("Avatar for {} not found anywhere.".format(username))
 
-    def read_img(self, pos, size):
+    def read_img_from_block(self, filename, pos, size):
         file_idx = pos >> 32
         fname = os.path.join(self.sfs_dir,
                 'avatar.block.' + '{:05d}'.format(file_idx))
-        # a 64-byte offset of each block file
-        start_pos = pos - file_idx * (2**32) + 64
+        # offset of each block file: 17 + len(path)
+        start_pos = pos - file_idx * (2**32) + 16 + len(filename) + 1
         try:
             with open(fname, 'rb') as f:
                 f.seek(start_pos)
@@ -95,11 +101,14 @@ class AvatarReader(object):
                     img[i,j] = (r, g, b)
             return Image.fromarray(img, mode="RGB")
 
-    def query_index(self, filename):
+    def query_index(self, avtid):
         conn = sqlite3.connect(self.avt_db)
-        cursor = conn.execute("select Offset,Size from Index_avatar where FileName='{}'".format(filename))
-        pos, size = cursor.fetchone()
-        return pos, size
+        cursor = conn.execute("select FileName,Offset,Size from Index_avatar")
+        candidates = []
+        for path, offset, size in cursor:
+            if avtid in path:
+                candidates.append((path, offset, size))
+        return candidates
 
 if __name__ == '__main__':
     import sys