Merge pull request ehForwarderBot#18 from RoyXiang/pr

[EWS]: Add an experimental flag: qr_target

Author: blueset

docs/EWS.md

@@ -104,3 +104,8 @@ eh_wechat_slave = {
   List out all other messages (more than first one) for link messages with more than one link from MPS accounts.
 * `max_quote_length` _(int)_  [Default: `-1`]  
   Maximum length of text for quotation messages. Set to 0 to fully disable quotation. Set to -1 to always quote full message.
+* `qr_target` _(string)_ [Default: `"console"`]
+  The target channel which should receive a QR code to login the Wechat account.
+  Available values:
+    * `"console"`: the command line interface
+    * `"master"`: the master channel

plugins/eh_wechat_slave.py

@@ -3,6 +3,7 @@
 import mimetypes
 import os
 import re
+import threading
 import time
 import html
 from binascii import crc32
@@ -100,6 +101,8 @@ class WeChatChannel(EFBChannel):
     supported_message_types = {MsgType.Text, MsgType.Sticker, MsgType.Image, MsgType.File, MsgType.Video, MsgType.Link}
     users = {}
     logger = logging.getLogger("plugins.%s.WeChatChannel" % channel_id)
+    qr_callback = ""
+    qr_uuid = ""
 
     SYSTEM_USERNAMES = ["filehelper", "newsapp", "fmessage", "weibo", "qqmail", "fmessage", "tmessage", "qmessage",
                          "qqsync", "floatbottle", "lbsapp", "shakeapp", "medianote", "qqfriend", "readerapp",
@@ -110,11 +113,14 @@ def __init__(self, queue):
         super().__init__(queue)
         self.itchat = itchat.new_instance()
         itchat.set_logging(showOnCmd=False)
-        self.itchat.auto_login(enableCmdQR=2,
-                               hotReload=True,
-                               statusStorageDir="storage/%s.pkl" % self.channel_id,
-                               exitCallback=self.exit_callback,
-                               qrCallback=self.console_qr_code)
+        self.itchat_msg_register()
+
+        qr_target = self._flag("qr_target", "console")
+        if qr_target == "master":
+            self.qr_callback = "master_qr_code"
+        else:
+            self.qr_callback = "console_qr_code"
+
         self.logger.info("EWS Inited!!!\n---")
 
     #
@@ -147,21 +153,54 @@ def console_qr_code(self, uuid, status, qrcode):
             QR += "\nIf you cannot read the QR code above, " \
                   "Please generate a QR code with any tool with the following URL:\n" \
                   "https://login.weixin.qq.com/l/" + uuid
-
         self.logger.critical(QR)
 
+    def master_qr_code(self, uuid, status, qrcode):
+        if uuid != self.qr_uuid:
+            path = os.path.join("storage", self.channel_id)
+            if not os.path.exists(path):
+                os.makedirs(path)
+            path = os.path.join(path, 'QR.jpg')
+            with open(path, 'wb') as f:
+                f.write(qrcode)
+            msg = EFBMsg(self)
+            msg.type = MsgType.Image
+            msg.source = MsgSource.System
+            msg.origin = {
+                'name': 'WeChat System Message',
+                'alias': 'WeChat System Message',
+                'uid': -1
+            }
+            msg.text = '[%s] Scan this QR code to login your account' % self.channel_id
+            msg.path = path
+            msg.file = open(path, 'rb')
+            msg.mime = 'image/jpeg'
+            self.queue.put(msg)
+            self.qr_uuid = uuid
+
     def exit_callback(self):
         msg = EFBMsg(self)
-        msg.type = MsgType.Text
         msg.source = MsgSource.System
         msg.origin = {
             'name': 'WeChat System Message',
             'alias': 'WeChat System Message',
             'uid': -1
         }
-
         msg.text = "WeChat server logged out the user."
-
+        if self.qr_callback == "console_qr_code":
+            msg.type = MsgType.Text
+        else:
+            msg.type = MsgType.Command
+            msg.attributes = {
+                "commands": [
+                    {
+                        "name": "Login again",
+                        "callable": "start_polling_thread",
+                        "args": [],
+                        "kwargs": {}
+                    }
+                ]
+            }
         self.queue.put(msg)
 
     def get_uid(self, UserName=None, NickName=None, alias=None, Uin=None):
@@ -301,8 +340,22 @@ def search_user(self, UserName=None, uid=None, uin=None, name=None, ActualUserNa
         return result
 
     def poll(self):
+        self.itchat.auto_login(enableCmdQR=2,
+                               hotReload=True,
+                               statusStorageDir="storage/%s.pkl" % self.channel_id,
+                               exitCallback=self.exit_callback,
+                               qrCallback=getattr(self, self.qr_callback))
         self.usersdata = self.itchat.get_friends(True) + self.itchat.get_chatrooms()
 
+        while self.itchat.alive and not self.stop_polling:
+            self.itchat.configured_reply()
+
+        if self.itchat.useHotReload:
+            self.itchat.dump_login_status()
+
+        self.logger.debug("%s (%s) gracefully stopped.", self.channel_name, self.channel_id)
+
+    def itchat_msg_register(self):
         @self.itchat.msg_register(['Text'], isFriendChat=True, isMpChat=True)
         def wcText(msg):
             self.textMsg(msg)
@@ -388,14 +441,6 @@ def wcSystemGroup(msg):
         def wcSysLog(msg):
             self.logger.debug("WeChat \"System\" message:\n%s", repr(msg))
 
-        while self.itchat.alive and not self.stop_polling:
-            self.itchat.configured_reply()
-
-        if self.itchat.useHotReload:
-            self.itchat.dump_login_status()
-
-        self.logger.debug("%s (%s) gracefully stopped.", self.channel_name, self.channel_id)
-
     @incomeMsgMeta
     def textMsg(self, msg, isGroupChat=False):
         self.logger.info("TextMsg!!!\n---")
@@ -772,6 +817,11 @@ def uin_rate(self, param=""):
 
     # Command functions
 
+    def start_polling_thread(self):
+        thread = threading.Thread(target=self.poll)
+        thread.start()
+        return "Reconnecting..."
+
     def add_friend(self, UserName=None, status=2, ticket="", UserInfo={}):
         if not UserName:
             return "Username is empty. (UE02)"
@@ -828,4 +878,4 @@ def _flag(self, key, value):
         Returns:
             Value for the flag.
         """
-        return getattr(config, self.channel_id, dict()).get('flags', dict()).get(key, value)
+        return getattr(config, self.channel_id, dict()).get('flags', dict()).get(key, value)