Pārlūkot izejas kodu

fix old avatar storage

Yuxin Wu 5 gadi atpakaļ
vecāks
revīzija
eeba4b7e6c
4 mainītis faili ar 69 papildinājumiem un 52 dzēšanām
  1. 1 0
      .gitignore
  2. 29 19
      android-interact.sh
  3. 0 3
      plot-num-msg-by-time.py
  4. 39 30
      wechat/avatar.py

+ 1 - 0
.gitignore

@@ -3,6 +3,7 @@
 output/
 resource
 *.db
+*.db.decoded
 *2015*
 *2017*
 2020*

+ 29 - 19
android-interact.sh

@@ -1,6 +1,5 @@
 #!/bin/bash
 # File: android-interact.sh
-# Date: Wed Nov 29 02:19:00 2017 -0800
 
 PROG_NAME=`python -c "import os, sys; print(os.path.realpath(sys.argv[1]))" "$0"`
 PROG_DIR=`dirname "$PROG_NAME"`
@@ -29,25 +28,36 @@ if [[ $1 == "db" || $1 == "res" ]]; then
 	echo "Found $numUser user(s). User chosen: $chooseUser"
 
 	if [[ $1 == "res" ]]; then
-		mkdir -p resource; cd resource
-		echo "Pulling resources... "
-		for d in avatar image2 voice2 emoji video sfs; do
-			adb shell "cd $RES_DIR/$chooseUser &&
-								 tar czf - $d 2>/dev/null | base64" |
-					base64 -di | tar xzf -
+		mkdir -p resource
+    (
+      cd resource || exit
+      echo "Pulling resources... "
+      for d in avatar image2 voice2 emoji video sfs; do
+        echo "Trying to download $RES_DIR/$chooseUser/$d with busybox ..."
+        adb shell "cd $RES_DIR/$chooseUser &&
+                   busybox tar czf - $d 2>/dev/null | busybox base64" |
+            base64 -di | tar xzf -
+        [[ -d $d ]] && continue
 
-			# Old Slow Way:
-			# mkdir -p $d; cd $d
-			# adb pull "$RES_DIR/$chooseUser/$d"
-			# cd ..
-			[[ -d $d ]] || {
-				>&2 echo "Failed to download resource directory: $RES_DIR/$chooseUser/$d"
-				exit 1
-			}
-		done
-		cd ..
-		echo "Resource pulled at ./resource"
-		echo "Total size: $(du -sh resource | cut -f1)"
+        echo "Trying to download $RES_DIR/$chooseUser/$d with tar & base64 ..."
+        adb shell "cd $RES_DIR/$chooseUser &&
+                   tar czf - $d 2>/dev/null | base64" | base64 -di | tar xzf -
+        [[ -d $d ]] && continue
+
+        echo "Trying to download $RES_DIR/$chooseUser/$d with adb pull (slow) ..."
+        mkdir -p $d
+        (
+          cd $d || exit
+          adb pull "$RES_DIR/$chooseUser/$d"
+        )
+
+        [[ -d $d ]] || {
+          echo "Failed to download $RES_DIR/$chooseUser/$d"
+        }
+      done
+      echo "Resource pulled at ./resource"
+      echo "Total size: $(du -sh | cut -f1)"
+    )
 	else
 		echo "Pulling database and avatar index file..."
 		adb pull $MM_DIR/MicroMsg/$chooseUser/EnMicroMsg.db

+ 0 - 3
plot-num-msg-by-time.py

@@ -1,8 +1,5 @@
 #!/usr/bin/env python3
 # -*- coding: UTF-8 -*-
-# File: plot-num-msg-by-time.py
-# Date: Wed Mar 25 17:44:39 2015 +0800
-# Author: Yuxin Wu
 
 from wechat.parser import WeChatDBParser
 from wechat.common.textutil import ensure_unicode

+ 39 - 30
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