diff --git a/dump_html.py b/dump_html.py index f0daa79..ecc3de6 100755 --- a/dump_html.py +++ b/dump_html.py @@ -1,24 +1,29 @@ #!/usr/bin/env python2 # -*- coding: UTF-8 -*- # File: dump_html.py -# Date: Sat Nov 22 23:51:45 2014 +0800 +# Date: Sun Nov 23 16:37:48 2014 +0800 # Author: Yuxin Wu import sys -if len(sys.argv) != 3: - sys.exit("Usage: {0} ".format(sys.argv[0])) +if len(sys.argv) != 5: + sys.exit("Usage: {0} ".format(sys.argv[0])) from lib.utils import ensure_unicode from lib.parser import WeChatDBParser +from lib.res import Resource from lib.render import HTMLRender db_file = sys.argv[1] -name = ensure_unicode(sys.argv[2]) +res = Resource(sys.argv[2]) +name = ensure_unicode(sys.argv[3]) +output_file = sys.argv[4] parser = WeChatDBParser(db_file) parser.parse() msgs = parser.msgs_by_talker[name] -render = HTMLRender() +render = HTMLRender(parser, res) html = render.render_msgs(msgs).encode('utf-8') -print html + +with open(output_file, 'w') as f: + print >> f, html diff --git a/lib/avatar.py b/lib/avatar.py new file mode 100644 index 0000000..7c09c2c --- /dev/null +++ b/lib/avatar.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python2 +# -*- coding: UTF-8 -*- +# File: avatar.py +# Date: Sun Nov 23 16:28:00 2014 +0800 +# Author: Yuxin Wu + +import os +import md5 +import numpy as np + +from .utils import ensure_bin_str + +class AvatarReader(object): + def __init__(self, res_dir): + self.avt_dir = os.path.join(res_dir, 'avatar') + assert os.path.isdir(self.avt_dir), "No such directory {}".format(self.avt_dir) + + @staticmethod + def get_filename(username): + m = md5.new() + m.update(username) + return m.hexdigest() + + def get_avatar(self, username): + """ username: `username` field in db.rcontact""" + username = ensure_bin_str(username) + filename = AvatarReader.get_filename(username) + dir1, dir2 = filename[:2], filename[2:4] + filename = os.path.join(self.avt_dir, dir1, dir2, + "user_{}.png.bm".format(filename)) + if not os.path.isfile(filename): + # avatar not found! + return None + else: + return AvatarReader.read_bm(filename) + + @staticmethod + def read_bm(fname): + size = (96, 96, 3) + img = np.zeros(size, dtype='uint8') + with open(fname, 'rb') as f: + 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 diff --git a/lib/parser.py b/lib/parser.py index efdd897..d04a26a 100644 --- a/lib/parser.py +++ b/lib/parser.py @@ -1,7 +1,7 @@ #!/usr/bin/env python2 # -*- coding: UTF-8 -*- # File: parser.py -# Date: Sat Nov 22 23:23:44 2014 +0800 +# Date: Sun Nov 23 16:33:20 2014 +0800 # Author: Yuxin Wu import sqlite3 @@ -18,7 +18,6 @@ addr_upload2 chatroom message rcontact - """ class WeChatDBParser(object): @@ -42,7 +41,10 @@ SELECT username,conRemark,nickname FROM rcontact else: self.contacts[username] = ensure_unicode(nickname) + print "Got {} contacts.".format(len(self.contacts)) + def _parse_msg(self): + msgs_tot_cnt = 0 db_msgs = self.cc.execute( """ SELECT {} FROM message @@ -57,6 +59,14 @@ SELECT {} FROM message for k, v in self.msgs_by_talker.iteritems(): for msg in v: msg.talker_name = ensure_unicode(k) + msgs_tot_cnt += len(v) + print "Got {} messages in total.".format(msgs_tot_cnt) + + def _parse_userinfo(self): + userinfo_q = self.cc.execute(""" SELECT id, value FROM userinfo """) + userinfo = dict(userinfo_q) + self.username = userinfo[2] + print "Your username is: {}".format(self.username) def _find_msg_by_type(self, msgs=None): ret = [] @@ -68,5 +78,6 @@ SELECT {} FROM message return sorted(ret) def parse(self): + self._parse_userinfo() self._parse_contact() self._parse_msg() diff --git a/lib/render.py b/lib/render.py index 171948d..93b8ed9 100755 --- a/lib/render.py +++ b/lib/render.py @@ -1,7 +1,7 @@ #!/usr/bin/env python2 # -*- coding: UTF-8 -*- # File: render.py -# Date: Sat Nov 22 23:24:10 2014 +0800 +# Date: Sun Nov 23 16:28:28 2014 +0800 # Author: Yuxin Wu import os @@ -22,14 +22,20 @@ TEMPLATES = dict([(k, open(os.path.join( LIB_PATH, 'static/{}.html'.format(v))).read()) for k, v in TEMPLATES_FILES.iteritems()]) - - class HTMLRender(object): - def __init__(self, res=None): + def __init__(self, parser, res=None): self.css = ensure_unicode(css_compress(open(CSS_FILE).read())) self.html = ensure_unicode(open(HTML_FILE).read()) + self.parser = parser self.res = res + def get_avatar_pair(self, username): + if self.res is None: + return ("", "") + avt1 = self.res.get_avatar(self.parser.username) + avt2 = self.res.get_avatar(username) + return (avt1, avt2) + def render_msg(self, msg): """ render a message, return the block""" # TODO @@ -42,9 +48,13 @@ class HTMLRender(object): content=msg.msg_str()) def render_msgs(self, msgs): + talker_name = msgs[0].talker + avatars = self.get_avatar_pair(talker_name) blocks = [self.render_msg(m) for m in msgs] return self.html.format(style=self.css, talker=msgs[0].talker_name, - messages=u''.join(blocks)) + messages=u''.join(blocks), + avatar_me=avatars[0], + avatar_you=avatars[1]) if __name__ == '__main__': r = HTMLRender() diff --git a/lib/res.py b/lib/res.py index b9694bd..c3f722c 100644 --- a/lib/res.py +++ b/lib/res.py @@ -1,11 +1,16 @@ #!/usr/bin/env python2 # -*- coding: UTF-8 -*- -# File: Res.py -# Date: Sat Nov 22 20:54:25 2014 +0800 +# File: res.py +# Date: Sun Nov 23 15:59:48 2014 +0800 # Author: Yuxin Wu import glob import os +import Image +import cStringIO +import base64 + +from lib.avatar import AvatarReader VOICE_DIRNAME = 'voice2' @@ -14,6 +19,7 @@ class Resource(object): def __init__(self, res_dir): assert os.path.isdir(res_dir), "No such directory: {}".format(res_dir) self.res_dir = res_dir + self.avt_reader = AvatarReader(self.res_dir) self.init() def init(self): @@ -32,6 +38,16 @@ class Resource(object): "Error interpreting the protocol, this is a bug!" self.speak_data[key] = full_path + def get_avatar(self, username): + """ return base64 string""" + ret = self.avt_reader.get_avatar(username) + if ret is None: + return "" + im = Image.fromarray(ret) + buf = cStringIO.StringIO() + im.save(buf, 'JPEG', quality=50) + jpeg_str = buf.getvalue() + return base64.b64encode(jpeg_str) diff --git a/lib/static/template.html b/lib/static/template.html index cf80842..62c6d1a 100644 --- a/lib/static/template.html +++ b/lib/static/template.html @@ -10,14 +10,14 @@
-
+

{talker}

-
+
{messages}
@@ -26,5 +26,17 @@
+