From 77191e1885d364592c3be574fdd8658a1a400fd8 Mon Sep 17 00:00:00 2001 From: Yuxin Wu Date: Wed, 20 Jan 2016 14:11:40 -0500 Subject: [PATCH] group message with proper handling of user name --- README.md | 2 -- dump-html.py | 4 +-- dump-msg.py | 2 +- plot-num-msg-by-time.py | 4 +-- wechat/libchathelper.py | 4 +-- wechat/msg.py | 38 ++++++------------------- wechat/parser.py | 57 +++++++++++++++++++++++++++++-------- wechat/render.py | 18 ++++++------ wechat/static/TP_INDEX.html | 4 +-- 9 files changed, 71 insertions(+), 62 deletions(-) diff --git a/README.md b/README.md index c51acc7..a04cc47 100644 --- a/README.md +++ b/README.md @@ -83,11 +83,9 @@ Screenshots of generated html: ![byvoid](https://github.com/ppwwyyxx/wechat-dump/raw/master/screenshots/byvoid.jpg) ### TODO List -+ Parse group messages. Doesn't work very well for now. + Search by uid/username + Faster way to copy a directory from android (I don't know..). + Fix rare unhandled types: > 10000 and < 0 -+ Use alias name in group chat, instead of id + Better user experiences... see `grep 'TODO' wechat -R` + more easy-to-use for non-programmers (GUI?) diff --git a/dump-html.py b/dump-html.py index c0bc642..ee9c39a 100755 --- a/dump-html.py +++ b/dump-html.py @@ -23,9 +23,9 @@ parser = WeChatDBParser(db_file) res = Resource(resource_dir, avt_db) try: - msgs = parser.msgs_by_talker[name] + msgs = parser.msgs_by_chat[name] except: - sys.stderr.write(u"Valid Contacts: {}\n".format(u'\n'.join(parser.msgs_by_talker.keys()))) + sys.stderr.write(u"Valid Contacts: {}\n".format(u'\n'.join(parser.msgs_by_chat.keys()))) sys.stderr.write(u"Couldn't find that contact {}.".format(name)); sys.exit(1) diff --git a/dump-msg.py b/dump-msg.py index cd006ac..6233867 100755 --- a/dump-msg.py +++ b/dump-msg.py @@ -22,7 +22,7 @@ if not os.path.isdir(output_dir): parser = WeChatDBParser(db_file) -for name, msgs in parser.msgs_by_talker.iteritems(): +for name, msgs in parser.msgs_by_chat.iteritems(): print u"Writing msgs for {}".format(name) safe_name = safe_filename(name) with open(os.path.join(output_dir, safe_name + '.txt'), 'w') as f: diff --git a/plot-num-msg-by-time.py b/plot-num-msg-by-time.py index 9ef0cd3..0209a2f 100755 --- a/plot-num-msg-by-time.py +++ b/plot-num-msg-by-time.py @@ -1,6 +1,6 @@ #!/usr/bin/env python2 # -*- coding: UTF-8 -*- -# File: plot_num_msg_by_time.py +# File: plot-num-msg-by-time.py # Date: Wed Mar 25 17:44:39 2015 +0800 # Author: Yuxin Wu @@ -20,7 +20,7 @@ name = ensure_unicode(sys.argv[2]) every_k_days = 2 parser = WeChatDBParser(db_file) -msgs = parser.msgs_by_talker[name] +msgs = parser.msgs_by_chat[name] times = [x.createTime for x in msgs] start_time = times[0] diffs = [(x - start_time).days for x in times] diff --git a/wechat/libchathelper.py b/wechat/libchathelper.py index febce1c..cb72c53 100644 --- a/wechat/libchathelper.py +++ b/wechat/libchathelper.py @@ -78,9 +78,9 @@ class LibChatHelper(object): return json.dumps(ret) def _convert_msg(self, msg): - sender = 'me' if msg.isSend else msg.get_msg_talker_id() + sender = 'me' if msg.isSend else msg.talker chatroom = msg.get_chatroom() - text = msg.content_no_first_line if msg.type == TYPE_MSG else '' + text = msg.content if msg.type == TYPE_MSG else '' img, format = self._get_image(msg) if img: # TODO don't use b64, directly return image content diff --git a/wechat/msg.py b/wechat/msg.py index 50e5a0a..586c194 100644 --- a/wechat/msg.py +++ b/wechat/msg.py @@ -21,7 +21,6 @@ TYPE_APP_MSG = 16777265 _KNOWN_TYPES = [eval(k) for k in dir() if k.startswith('TYPE_')] import re -from datetime import datetime from pyquery import PyQuery import logging logger = logging.getLogger(__name__) @@ -30,7 +29,6 @@ from common.textutil import ensure_unicode class WeChatMsg(object): - FIELDS = ["msgSvrId","type","isSend","createTime","talker","content","imgPath"] @staticmethod def filter_type(tp): @@ -38,21 +36,13 @@ class WeChatMsg(object): return True return False - def __init__(self, row): - """ row: a tuple corresponding to FIELDS""" - assert len(row) == len(WeChatMsg.FIELDS) - for f, v in zip(WeChatMsg.FIELDS, row): - setattr(self, f, v) + def __init__(self, values): + for k, v in values.iteritems(): + setattr(self, k, v) if self.type not in _KNOWN_TYPES: logger.warn("Unhandled message type: {}".format(self.type)) # only to supress repeated warning: _KNOWN_TYPES.append(self.type) - self.createTime = datetime.fromtimestamp(self.createTime / 1000) - self.talker_name = None - if self.content: - self.content = ensure_unicode(self.content) - else: - self.content = u"" def msg_str(self): if self.type == TYPE_LOCATION: @@ -97,29 +87,22 @@ class WeChatMsg(object): return "LOCATION SHARING" elif self.type == TYPE_EMOJI: # TODO add emoji name - return self.content_no_first_line + return self.content else: # TODO replace smiley with text - return self.content_no_first_line - - @property - def content_no_first_line(self): - if not self.is_chatroom(): return self.content - return self.content[self.content.find('\n')+1:] @property def content_xml_ready(self): # remove xml headers to avoid possible errors it may create header = re.compile(r'<\?.*\?>') - msg = header.sub("", self.content_no_first_line) + msg = header.sub("", self.content) return msg def __repr__(self): ret = u"{}|{}:{}:{}".format( self.type, - (self.talker if not self.talker_name else self.talker_name) \ - if not self.isSend else 'me', + self.talker if not self.isSend else 'me', self.createTime, ensure_unicode(self.msg_str())).encode('utf-8') if self.imgPath: @@ -132,16 +115,11 @@ class WeChatMsg(object): return self.createTime < r.createTime def is_chatroom(self): - return self.talker.endswith('@chatroom') - - def get_msg_talker_id(self): - if not self.is_chatroom(): - return self.talker - return self.content[:self.content.find(':')] + return self.talker != self.chat def get_chatroom(self): if self.is_chatroom(): - return self.talker[:-9] + return self.chat else: return '' diff --git a/wechat/parser.py b/wechat/parser.py index d87e562..0b755d0 100644 --- a/wechat/parser.py +++ b/wechat/parser.py @@ -7,6 +7,7 @@ import sqlite3 from collections import defaultdict import itertools +from datetime import datetime import logging logger = logging.getLogger(__name__) @@ -23,13 +24,15 @@ rcontact """ class WeChatDBParser(object): + FIELDS = ["msgSvrId","type","isSend","createTime","talker","content","imgPath"] + def __init__(self, db_fname): """ db_fname: EnMicroMsg.db""" self.db_fname = db_fname self.db_conn = sqlite3.connect(self.db_fname) self.cc = self.db_conn.cursor() self.contacts = {} - self.msgs_by_talker = defaultdict(list) + self.msgs_by_chat = defaultdict(list) self.emojis = {} self.internal_emojis = {} self._parse() @@ -46,6 +49,7 @@ SELECT username,conRemark,nickname FROM rcontact else: self.contacts[username] = ensure_unicode(nickname) + self.contacts_rev = {v: k for k, v in self.contacts.iteritems()} logger.info("Found {} contacts.".format(len(self.contacts))) def _parse_msg(self): @@ -53,19 +57,18 @@ SELECT username,conRemark,nickname FROM rcontact db_msgs = self.cc.execute( """ SELECT {} FROM message -""".format(','.join(WeChatMsg.FIELDS))) +""".format(','.join(WeChatDBParser.FIELDS))) for row in db_msgs: - msg = WeChatMsg(row) + values = self._parse_row(row) + if not values: + continue + msg = WeChatMsg(values) + # TODO keep system message? if not WeChatMsg.filter_type(msg.type): - self.msgs_by_talker[msg.talker].append(msg) + self.msgs_by_chat[msg.chat].append(msg) - # It's possible that messages are kept in database after contacts been deleted - # TODO handle this with a random contact name - self.msgs_by_talker = {self.contacts[k]: sorted(v, key=lambda x: x.createTime) - for k, v in self.msgs_by_talker.iteritems() if k in self.contacts} - for k, v in self.msgs_by_talker.iteritems(): - for msg in v: - msg.talker_name = ensure_unicode(k) + for k, v in self.msgs_by_chat.iteritems(): + self.msgs_by_chat[k] = sorted(v, key=lambda x: x.createTime) msgs_tot_cnt += len(v) logger.info("Found {} message records.".format(msgs_tot_cnt)) @@ -84,7 +87,7 @@ SELECT {} FROM message def _find_msg_by_type(self, msgs=None): ret = [] if msgs is None: - msgs = itertools.chain.from_iterable(self.msgs_by_talker.itervalues()) + msgs = itertools.chain.from_iterable(self.msgs_by_chat.itervalues()) for msg in msgs: if msg.type == 34: ret.append(msg) @@ -114,3 +117,33 @@ SELECT {} FROM message self._parse_msg() self._parse_imginfo() self._parse_emoji() + + # process the values in a row + def _parse_row(self, row): + values = dict(zip(WeChatDBParser.FIELDS, row)) + if values['content']: + values['content'] = ensure_unicode(values['content']) + else: + values['content'] = u'' + values['createTime'] = datetime.fromtimestamp(values['createTime']/ 1000) + values['chat'] = values['talker'] + try: + if values['chat'].endswith('@chatroom'): + values['chat'] = self.contacts[values['chat']] + content = values['content'] + talker = content[:content.find(':')] + try: + values['talker'] = self.contacts[talker] + values['content'] = content[content.find('\n') + 1:] + except KeyError: + # system messages have no talker + values['talker'] = u'' + else: + tk_id = values['talker'] + values['chat'] = self.contacts[tk_id] + values['talker'] = self.contacts[tk_id] + except KeyError: + # It's possible that messages are kept in database after contacts been deleted + logger.warn("Unknown contact, probably deleted: {}".format(tk_id)) + return None + return values diff --git a/wechat/render.py b/wechat/render.py index 673e609..bfc2718 100644 --- a/wechat/render.py +++ b/wechat/render.py @@ -92,7 +92,7 @@ class HTMLRender(object): def render_msg(self, msg): """ render a message, return the html block""" # TODO for chatroom, add nickname on avatar - sender = u'you ' + msg.get_msg_talker_id() if not msg.isSend else 'me' + sender = u'you ' + msg.talker if not msg.isSend else 'me' format_dict = {'sender_label': sender, 'time': msg.createTime } def fallback(): @@ -174,7 +174,7 @@ class HTMLRender(object): # string operation is extremely slow return self.html.format(extra_css=self.all_css, extra_js=self.all_js, - talker=msgs[0].talker_name, + chat=msgs[0].chat, messages=u''.join(blocks) ) @@ -184,25 +184,25 @@ class HTMLRender(object): css = avatar_tpl.format(name='me', avatar=my_avatar) for talker in talkers: - avatar = self.res.get_avatar(talker) + avatar = self.res.get_avatar(self.parser.contacts_rev[talker]) css += avatar_tpl.format(name=talker, avatar=avatar) self.css_string.append(css) def render_msgs(self, msgs): - """ render msgs of one friend, return a list of html""" - talker_id = msgs[0].talker + """ render msgs of one chat, return a list of html""" + chat = msgs[0].chat if msgs[0].is_chatroom(): talkers = set() for msg in msgs: - talkers.add(msg.get_msg_talker_id()) + talkers.add(msg.talker) else: - talkers = set([talker_id]) + talkers = set([chat]) self.prepare_avatar_css(talkers) self.res.cache_voice_mp3(msgs) - logger.info(u"Rendering {} messages of {}({})".format( - len(msgs), self.parser.contacts[talker_id], talker_id)) + logger.info(u"Rendering {} messages of {}".format( + len(msgs), chat)) self.prgs = ProgressReporter("Render", total=len(msgs)) slice_by_size = MessageSlicerBySize().slice(msgs) diff --git a/wechat/static/TP_INDEX.html b/wechat/static/TP_INDEX.html index c429a49..a1559d2 100644 --- a/wechat/static/TP_INDEX.html +++ b/wechat/static/TP_INDEX.html @@ -2,7 +2,7 @@ - Chat with {talker} + Chat with {chat} {extra_css} {extra_js} @@ -12,7 +12,7 @@
-

{talker}

+

{chat}