hotreload.py 3.9KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. import pickle, os
  2. import logging
  3. import requests # type: ignore
  4. from ..config import VERSION
  5. from ..returnvalues import ReturnValue
  6. from ..storage import templates
  7. from .contact import update_local_chatrooms, update_local_friends
  8. from .messages import produce_msg
  9. logger = logging.getLogger('itchat')
  10. def load_hotreload(core):
  11. core.dump_login_status = dump_login_status
  12. core.load_login_status = load_login_status
  13. async def dump_login_status(self, fileDir=None):
  14. fileDir = fileDir or self.hotReloadDir
  15. try:
  16. with open(fileDir, 'w') as f:
  17. f.write('itchat - DELETE THIS')
  18. os.remove(fileDir)
  19. except:
  20. raise Exception('Incorrect fileDir')
  21. status = {
  22. 'version' : VERSION,
  23. 'loginInfo' : self.loginInfo,
  24. 'cookies' : self.s.cookies.get_dict(),
  25. 'storage' : self.storageClass.dumps()}
  26. with open(fileDir, 'wb') as f:
  27. pickle.dump(status, f)
  28. logger.debug('Dump login status for hot reload successfully.')
  29. async def load_login_status(self, fileDir,
  30. loginCallback=None, exitCallback=None):
  31. try:
  32. with open(fileDir, 'rb') as f:
  33. j = pickle.load(f)
  34. except Exception as e:
  35. logger.debug('No such file, loading login status failed.')
  36. return ReturnValue({'BaseResponse': {
  37. 'ErrMsg': 'No such file, loading login status failed.',
  38. 'Ret': -1002, }})
  39. if j.get('version', '') != VERSION:
  40. logger.debug(('you have updated itchat from %s to %s, ' +
  41. 'so cached status is ignored') % (
  42. j.get('version', 'old version'), VERSION))
  43. return ReturnValue({'BaseResponse': {
  44. 'ErrMsg': 'cached status ignored because of version',
  45. 'Ret': -1005, }})
  46. self.loginInfo = j['loginInfo']
  47. self.loginInfo['User'] = templates.User(self.loginInfo['User'])
  48. self.loginInfo['User'].core = self
  49. self.s.cookies = requests.utils.cookiejar_from_dict(j['cookies'])
  50. self.storageClass.loads(j['storage'])
  51. try:
  52. msgList, contactList = self.get_msg()
  53. except:
  54. msgList = contactList = None
  55. if (msgList or contactList) is None:
  56. self.logout()
  57. await load_last_login_status(self.s, j['cookies'])
  58. logger.debug('server refused, loading login status failed.')
  59. return ReturnValue({'BaseResponse': {
  60. 'ErrMsg': 'server refused, loading login status failed.',
  61. 'Ret': -1003, }})
  62. else:
  63. if contactList:
  64. for contact in contactList:
  65. if '@@' in contact['UserName']:
  66. update_local_chatrooms(self, [contact])
  67. else:
  68. update_local_friends(self, [contact])
  69. if msgList:
  70. msgList = produce_msg(self, msgList)
  71. for msg in msgList: self.msgList.put(msg)
  72. await self.start_receiving(exitCallback)
  73. logger.debug('loading login status succeeded.')
  74. if hasattr(loginCallback, '__call__'):
  75. await loginCallback(self.storageClass.userName)
  76. return ReturnValue({'BaseResponse': {
  77. 'ErrMsg': 'loading login status succeeded.',
  78. 'Ret': 0, }})
  79. async def load_last_login_status(session, cookiesDict):
  80. try:
  81. session.cookies = requests.utils.cookiejar_from_dict({
  82. 'webwxuvid': cookiesDict['webwxuvid'],
  83. 'webwx_auth_ticket': cookiesDict['webwx_auth_ticket'],
  84. 'login_frequency': '2',
  85. 'last_wxuin': cookiesDict['wxuin'],
  86. 'wxloadtime': cookiesDict['wxloadtime'] + '_expired',
  87. 'wxpluginkey': cookiesDict['wxloadtime'],
  88. 'wxuin': cookiesDict['wxuin'],
  89. 'mm_lang': 'zh_CN',
  90. 'MM_WX_NOTIFY_STATE': '1',
  91. 'MM_WX_SOUND_STATE': '1', })
  92. except:
  93. logger.info('Load status for push login failed, we may have experienced a cookies change.')
  94. logger.info('If you are using the newest version of itchat, you may report a bug.')