Enhanced pkl storage & add qrCallback & fix contact self bug

littlecodersh · littlecodersh · commit c30ae9a1bde9 · 2016-12-04T13:51:20.000+08:00
diff --git a/itchat/__init__.py b/itchat/__init__.py
@@ -1,8 +1,9 @@
-from .core import Core
 from . import content
+from .core import Core
+from .config import VERSION
 from .log import set_logging
 
-__version__ = '1.2.9'
+__version__ = VERSION
 
 instanceList = []
 
diff --git a/itchat/components/contact.py b/itchat/components/contact.py
@@ -118,8 +118,6 @@ def update_local_chatrooms(core, l):
         for member in chatroom['MemberList']:
             utils.emoji_formatter(member, 'NickName')
             utils.emoji_formatter(member, 'DisplayName')
-            if core.storageClass.userName == member['UserName']:
-                chatroom['self'] = member
         # update it to old chatrooms
         oldChatroom = utils.search_dict_list(
             core.chatroomList, 'UserName', chatroom['UserName'])
@@ -149,6 +147,10 @@ def update_local_chatrooms(core, l):
                 oldChatroom['OwnerUin'] == int(core.loginInfo['wxuin'])
         else:
             oldChatroom['isAdmin'] = None
+        #  - update self
+        newSelf = utils.search_dict_list(oldChatroom['MemberList'],
+            'UserName', core.storageClass.userName)
+        oldChatroom['self'] = newSelf or copy.deepcopy(core.loginInfo['User'])
     return {
         'Type'         : 'System',
         'Text'         : [chatroom['UserName'] for chatroom in l],
diff --git a/itchat/components/hotreload.py b/itchat/components/hotreload.py
@@ -3,6 +3,7 @@
 
 import requests
 
+from ..config import VERSION
 from ..returnvalues import ReturnValue
 from .contact import update_local_chatrooms
 from .messages import produce_msg
@@ -22,6 +23,7 @@ def dump_login_status(self, fileDir=None):
     except:
         raise Exception('Incorrect fileDir')
     status = {
+        'version'   : VERSION,
         'loginInfo' : self.loginInfo,
         'cookies'   : self.s.cookies.get_dict(),
         'storage'   : self.storageClass.dumps()}
@@ -40,6 +42,13 @@ def load_login_status(self, fileDir,
             'ErrMsg': 'No such file, loading login status failed.',
             'Ret': -1002, }})
 
+    if j.get('version', '') != VERSION:
+        logger.debug(('you have updated itchat from %s to %s, ' + 
+            'so cached status is ignored') % (
+            j.get('version', 'old version'), VERSION))
+        return ReturnValue({'BaseResponse': {
+            'ErrMsg': 'cached status ignored because of version',
+            'Ret': -1005, }})
     self.loginInfo = j['loginInfo']
     self.s.cookies = requests.utils.cookiejar_from_dict(j['cookies'])
     self.storageClass.loads(j['storage'])
diff --git a/itchat/components/login.py b/itchat/components/login.py
@@ -1,4 +1,4 @@
-import os, sys, time, re
+import os, sys, time, re, io
 import threading
 import json, xml.dom.minidom
 import copy, pickle, random
@@ -24,7 +24,7 @@ def load_login(core):
     core.get_msg           = get_msg
     core.logout            = logout
 
-def login(self, enableCmdQR=False, picDir=None,
+def login(self, enableCmdQR=False, picDir=None, qrCallback=None,
         loginCallback=None, exitCallback=None):
     if self.alive:
         logger.debug('itchat has already logged in.')
@@ -34,19 +34,28 @@ def login(self, enableCmdQR=False, picDir=None,
             logger.info('Getting uuid of QR code.')
             while not self.get_QRuuid(): time.sleep(1)
             logger.info('Downloading QR code.')
-            if self.get_QR(enableCmdQR=enableCmdQR, picDir=picDir):
+            qrStorage = self.get_QR(enableCmdQR=enableCmdQR,
+                picDir=picDir, qrCallback=qrCallback)
+            if qrStorage:
                 break
             elif 9 == getCount:
                 logger.info('Failed to get QR code, please restart the program.')
                 sys.exit()
         logger.info('Please scan the QR code to log in.')
-        status = self.check_login()
-        if status == '201':
-            logger.info('Please press confirm on your phone.')
-            while status == '201':
-                status = self.check_login()
-                time.sleep(1)
-        if status == '200': break
+        isLoggedIn = False
+        while not isLoggedIn:
+            status = self.check_login()
+            if hasattr(qrCallback, '__call__'):
+                qrCallback(uuid=self.uuid, status=status, qrcode=qrStorage.getvalue())
+            if status == '200':
+                isLoggedIn = True
+            elif status == '201':
+                if isLoggedIn is not None:
+                    logger.info('Please press confirm on your phone.')
+                    isLoggedIn = None
+            elif status != '408':
+                break
+        if isLoggedIn: break
         logger.info('Log in time out, reloading QR code')
     self.web_init()
     self.show_mobile_login()
@@ -72,39 +81,43 @@ def get_QRuuid(self):
         self.uuid = data.group(2)
         return self.uuid
 
-def get_QR(self, uuid=None, enableCmdQR=False, picDir=None):
+def get_QR(self, uuid=None, enableCmdQR=False, picDir=None, qrCallback=None):
+    uuid = uuid or self.uuid
+    picDir = picDir or config.DEFAULT_QR
+    url = '%s/qrcode/%s' % (config.BASE_URL, uuid)
+    headers = { 'User-Agent' : config.USER_AGENT }
     try:
-        uuid = uuid or self.uuid
-        picDir = picDir or config.DEFAULT_QR
-        url = '%s/qrcode/%s' % (config.BASE_URL, uuid)
-        headers = { 'User-Agent' : config.USER_AGENT }
         r = self.s.get(url, stream=True, headers=headers)
-        with open(picDir, 'wb') as f: f.write(r.content)
     except:
         return False
-    if enableCmdQR:
-        utils.print_cmd_qr(picDir, enableCmdQR = enableCmdQR)
+    qrStorage = io.BytesIO(r.content)
+    if hasattr(qrCallback, '__call__'):
+        qrCallback(uuid=uuid, status='0', qrcode=qrStorage.getvalue())
     else:
-        utils.print_qr(picDir)
-    return True
+        with open(picDir, 'wb') as f: f.write(r.content)
+        if enableCmdQR:
+            utils.print_cmd_qr(picDir, enableCmdQR=enableCmdQR)
+        else:
+            utils.print_qr(picDir)
+    return qrStorage
 
 def check_login(self, uuid=None):
     uuid = uuid or self.uuid
     url = '%s/cgi-bin/mmwebwx-bin/login' % config.BASE_URL
-    params = 'tip=1&uuid=%s&_=%s' % (uuid, int(time.time()))
+    localTime = int(time.time())
+    params = 'loginicon=true&uuid=%s&tip=0&r=%s&_=%s' % (
+        uuid, localTime / 1579, localTime)
     headers = { 'User-Agent' : config.USER_AGENT }
     r = self.s.get(url, params=params, headers=headers)
     regx = r'window.code=(\d+)'
     data = re.search(regx, r.text)
     if data and data.group(1) == '200':
         process_login_info(self, r.text)
         return '200'
-    elif data and data.group(1) == '201':
-        return '201'
-    elif data and data.group(1) == '408':
-        return '408'
+    elif data:
+        return data.group(1)
     else:
-        return '0'
+        return '400'
 
 def process_login_info(core, loginContent):
     ''' when finish login (scanning qrcode)
diff --git a/itchat/components/register.py b/itchat/components/register.py
@@ -15,16 +15,19 @@ def load_register(core):
     core.run              = run
 
 def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
-        enableCmdQR=False, picDir=None, loginCallback=None, exitCallback=None):
+        enableCmdQR=False, picDir=None, qrCallback=None,
+        loginCallback=None, exitCallback=None):
     self.useHotReload = hotReload
     if hotReload:
-        if self.load_login_status(statusStorageDir, loginCallback=loginCallback, exitCallback=exitCallback): return
-        self.login(enableCmdQR=enableCmdQR, picDir=picDir,
+        if self.load_login_status(statusStorageDir,
+                loginCallback=loginCallback, exitCallback=exitCallback):
+            return
+        self.login(enableCmdQR=enableCmdQR, picDir=picDir, qrCallback=qrCallback,
             loginCallback=loginCallback, exitCallback=exitCallback)
         self.dump_login_status(statusStorageDir)
         self.hotReloadDir = statusStorageDir
     else:
-        self.login(enableCmdQR=enableCmdQR, picDir=picDir,
+        self.login(enableCmdQR=enableCmdQR, picDir=picDir, qrCallback=qrCallback,
             loginCallback=loginCallback, exitCallback=exitCallback)
 
 def configured_reply(self):
diff --git a/itchat/config.py b/itchat/config.py
@@ -1,5 +1,6 @@
 import os, platform
 
+VERSION = '1.2.10'
 BASE_URL = 'https://login.weixin.qq.com'
 OS = platform.system() #Windows, Linux, Darwin
 DIR = os.getcwd()
diff --git a/itchat/core.py b/itchat/core.py
@@ -29,7 +29,7 @@ def __init__(self):
         self.functionDict = {'FriendChat': {}, 'GroupChat': {}, 'MpChat': {}}
         self.useHotReload, self.hotReloadDir = False, 'itchat.pkl'
         self.receivingRetryCount = 5
-    def login(self, enableCmdQR=False, picDir=None,
+    def login(self, enableCmdQR=False, picDir=None, qrCallback=None,
             loginCallback=None, exitCallback=None):
         ''' log in like web wechat does
             for log in
@@ -40,6 +40,7 @@ def login(self, enableCmdQR=False, picDir=None,
                 - enableCmdQR: show qrcode in command line
                     - integers can be used to fit strange char length
                 - picDir: place for storing qrcode
+                - qrCallback: method that should accept uuid, status, qrcode
                 - loginCallback: callback after successfully logged in
                     - if not set, screen is cleared and qrcode is deleted
                 - exitCallback: callback after logged out
@@ -66,12 +67,13 @@ def get_QRuuid(self):
             it is defined in components/login.py
         '''
         raise NotImplementedError()
-    def get_QR(self, uuid=None, enableCmdQR=False, picDir=None):
+    def get_QR(self, uuid=None, enableCmdQR=False, picDir=None, qrCallback=None):
         ''' download and show qrcode
             for options
                 - uuid: if uuid is not set, latest uuid you fetched will be used
                 - enableCmdQR: show qrcode in cmd
                 - picDir: where to store qrcode
+                - qrCallback: method that should accept uuid, status, qrcode
             it is defined in components/login.py
         '''
         raise NotImplementedError()
@@ -385,7 +387,7 @@ def load_login_status(self, fileDir,
         '''
         raise NotImplementedError()
     def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
-            enableCmdQR=False, picDir=None,
+            enableCmdQR=False, picDir=None, qrCallback=None,
             loginCallback=None, exitCallback=None):
         ''' log in like web wechat does
             for log in
@@ -402,6 +404,7 @@ def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
                     - if not set, screen is cleared and qrcode is deleted
                 - exitCallback: callback after logged out
                     - it contains calling of logout
+                - qrCallback: method that should accept uuid, status, qrcode
             for usage
                 ..code::python
 
