From 268544183d6a798dc7cfb4e935ab394f514d7702 Mon Sep 17 00:00:00 2001
From: Yuxin Wu <ppwwyyxxc@gmail.com>
Date: Thu, 21 Apr 2016 18:12:57 -0400
Subject: [PATCH] update script with argparse

---
 dump-html.py     | 59 ++++++++++++++++++++++++++++--------------------
 list-chats.py    |  5 +++-
 wechat/avatar.py | 16 ++++++++-----
 3 files changed, 49 insertions(+), 31 deletions(-)

diff --git a/dump-html.py b/dump-html.py
index 1273e29..e5b442a 100755
--- a/dump-html.py
+++ b/dump-html.py
@@ -5,37 +5,48 @@
 # Author: Yuxin Wu <ppwwyyxxc@gmail.com>
 
 import sys
-if len(sys.argv) != 6:
-    sys.exit("Usage: {0} <path to decoded_database.db> <path to avatar.index> <path to resource> <name> <output html>".format(sys.argv[0]))
+import argparse
 
 from common.textutil import ensure_unicode
 from wechat.parser import WeChatDBParser
 from wechat.res import Resource
 from wechat.render import HTMLRender
 
-db_file = sys.argv[1]
-avt_db = sys.argv[2]
-resource_dir = sys.argv[3]
-name = ensure_unicode(sys.argv[4])
-output_file = sys.argv[5]
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('name', help='name of contact')
+    parser.add_argument('--output', help='output html file', default='output.html')
+    parser.add_argument('--db', default='decoded.db', help='path to decoded database')
+    parser.add_argument('--avt', default='avatar.index', help='path to avatar.index file')
+    parser.add_argument('--res', default='resource', help='reseource directory')
+    args = parser.parse_args()
+    return args
 
-parser = WeChatDBParser(db_file)
-res = Resource(resource_dir, avt_db)
+if __name__ == '__main__':
+    args = get_args()
 
-try:
-    msgs = parser.msgs_by_chat[name]
-except:
-    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)
+    name = ensure_unicode(args.name)
+    output_file = args.output
 
-render = HTMLRender(parser, res)
-htmls = render.render_msgs(msgs)
+    parser = WeChatDBParser(args.db)
+    res = Resource(args.res, args.avt)
 
-if len(htmls) == 1:
-    with open(output_file, 'w') as f:
-        print >> f, htmls[0].encode('utf-8')
-else:
-    for idx, html in enumerate(htmls):
-        with open(output_file + '.{}'.format(idx), 'w') as f:
-            print >> f, html.encode('utf-8')
+    if name and name in parser.msgs_by_chat:
+        msgs = parser.msgs_by_chat[name]
+    else:
+        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)
+    print "Number of Messages: ", len(msgs)
+    assert len(msgs) > 0
+
+    render = HTMLRender(parser, res)
+    htmls = render.render_msgs(msgs)
+
+    if len(htmls) == 1:
+        with open(output_file, 'w') as f:
+            print >> f, htmls[0].encode('utf-8')
+    else:
+        for idx, html in enumerate(htmls):
+            with open(output_file + '.{}'.format(idx), 'w') as f:
+                print >> f, html.encode('utf-8')
diff --git a/list-chats.py b/list-chats.py
index a64a434..eeea944 100755
--- a/list-chats.py
+++ b/list-chats.py
@@ -5,10 +5,13 @@
 
 from wechat.parser import WeChatDBParser
 import sys
+if len(sys.argv) != 2:
+    print "Usage: {} db_file".format(sys.argv[0])
+    sys.exit(1)
 
 db_file = sys.argv[1]
 
 parser = WeChatDBParser(db_file)
 chats = parser.msgs_by_chat.keys()
 for k in chats:
-    print k.encode('utf-8'), parser.contacts_rev[k].encode('utf-8')
+    print k.encode('utf-8'),  '\t', parser.contacts_rev[k].encode('utf-8')
diff --git a/wechat/avatar.py b/wechat/avatar.py
index 3eb4618..d84b6e6 100644
--- a/wechat/avatar.py
+++ b/wechat/avatar.py
@@ -30,10 +30,10 @@ class AvatarReader(object):
         try:
             try:
                 pos, size = self.query_index(filename)
+                return self.read_img(pos, size)
             except TypeError:
                 logger.warn("Avatar for {} not found in avatar database.".format(username))
                 return None
-            return self.read_img(pos, size)
         except Exception as e:
             raise
             print e
@@ -47,11 +47,15 @@ class AvatarReader(object):
                 'avatar.block.' + '{:05d}'.format(file_idx))
         # a 64-byte offset of each block file
         start_pos = pos - file_idx * (2**32) + 64
-        with open(fname, 'rb') as f:
-            f.seek(start_pos)
-            data = f.read(size)
-            im = Image.open(cStringIO.StringIO(data))
-            return im
+        try:
+            with open(fname, 'rb') as f:
+                f.seek(start_pos)
+                data = f.read(size)
+                im = Image.open(cStringIO.StringIO(data))
+                return im
+        except IOError as e:
+            logger.warn("Cannot read avatar from {}: {}".format(fname, str(e)))
+            return None
 
     def query_index(self, filename):
         conn = sqlite3.connect(self.avt_db)