ZSCGR/wechat-dump
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

check for two avatar storage

Browse Source
This commit is contained in:
Yuxin Wu
2020-06-28 00:33:07 -07:00
parent eeba4b7e6c
commit 1f1a28e605
1 changed files with 41 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wechat/avatar.py
+41 -33
View File
@@ -26,51 +26,59 @@ class AvatarReader(object):
# new location of avatar, see #50
self.avt_dir = os.path.join(res_dir, 'avatar')
if not os.path.isdir(self.avt_dir) or len(os.listdir(self.avt_dir)) == 0:
self.avt_dir = None
self.avt_db = avt_db
self._use_avt = True
if os.path.isdir(self.avt_dir) and len(os.listdir(self.avt_dir)):
# 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*')):
# has sfs/avatar*
pass
else:
if self.avt_db is not None:
if len(glob.glob(os.path.join(self.sfs_dir, 'avatar*'))) == 0:
# has sfs/avatar*
self.avt_db = None
if self.avt_dir is None and self.avt_db is None:
logger.warn("Cannot find avatar storage. Will not use avatar!")
self._use_avt = False
def get_avatar_from_avtdb(self, avtid):
try:
candidates = self._search_avt_db(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 Exception:
pass
def get_avatar_from_avtdir(self, avtid):
dir1, dir2 = avtid[:2], avtid[2:4]
candidates = glob.glob(os.path.join(self.avt_dir, dir1, dir2, f"*{avtid}*"))
candidates = sorted(set(candidates), key=_filename_priority, reverse=True)
for cand in candidates:
try:
if cand.endswith(".bm"):
return self.read_bm_file(cand)
else:
return Image.open(cand)
except Exception:
logger.exception("")
pass
def get_avatar(self, username):
""" username: `username` field in db.rcontact"""
if not self._use_avt:
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)
if self.avt_db is not None:
ret = self.get_avatar_from_avtdb(avtid)
if ret is not None:
return ret
candidates = sorted(set(candidates), key=_filename_priority, reverse=True)
for cand in candidates:
try:
if os.path.exists(cand):
if cand.endswith(".bm"):
return self.read_bm_file(cand)
else:
return Image.open(cand)
except Exception:
logger.exception("")
pass
if self.avt_dir is not None:
ret = self.get_avatar_from_avtdir(avtid)
if ret is not None:
return ret
logger.warning("Avatar for {} not found anywhere.".format(username))
def read_img_from_block(self, filename, pos, size):
@@ -101,7 +109,7 @@ class AvatarReader(object):
img[i,j] = (r, g, b)
return Image.fromarray(img, mode="RGB")
def query_index(self, avtid):
def _search_avt_db(self, avtid):
conn = sqlite3.connect(self.avt_db)
cursor = conn.execute("select FileName,Offset,Size from Index_avatar")
candidates = []