From 4cbf46fd4d0b5f175e15a55f71398bf629172bfd Mon Sep 17 00:00:00 2001
From: lanvent <lanvent@qq.com>
Date: Thu, 20 Apr 2023 01:03:04 +0800
Subject: [PATCH] feat: add support for wechatcom channel

---
 app.py                                 |   9 ++-
 channel/channel_factory.py             |   4 +
 channel/wechat/wechat_channel.py       |   4 +-
 channel/wechat/wechat_message.py       |   2 +-
 channel/wechatcom/wechatcom_channel.py | 105 +++++++++++++++++++++++++
 channel/wechatcom/wechatcom_message.py |  54 +++++++++++++
 channel/wechatmp/wechatmp_channel.py   |   6 +-
 config.py                              |   7 ++
 requirements-optional.txt              |   3 +-
 9 files changed, 187 insertions(+), 7 deletions(-)
 create mode 100644 channel/wechatcom/wechatcom_channel.py
 create mode 100644 channel/wechatcom/wechatcom_message.py

diff --git a/app.py b/app.py
index 637b6e4..1a1bb5e 100644
--- a/app.py
+++ b/app.py
@@ -43,7 +43,14 @@ def run():
             # os.environ['WECHATY_PUPPET_SERVICE_ENDPOINT'] = '127.0.0.1:9001'
 
         channel = channel_factory.create_channel(channel_name)
-        if channel_name in ["wx", "wxy", "terminal", "wechatmp", "wechatmp_service"]:
+        if channel_name in [
+            "wx",
+            "wxy",
+            "terminal",
+            "wechatmp",
+            "wechatmp_service",
+            "wechatcom",
+        ]:
             PluginManager().load_plugins()
 
         # startup channel
diff --git a/channel/channel_factory.py b/channel/channel_factory.py
index ebd9732..fed3234 100644
--- a/channel/channel_factory.py
+++ b/channel/channel_factory.py
@@ -29,4 +29,8 @@ def create_channel(channel_type):
         from channel.wechatmp.wechatmp_channel import WechatMPChannel
 
         return WechatMPChannel(passive_reply=False)
+    elif channel_type == "wechatcom":
+        from channel.wechatcom.wechatcom_channel import WechatComChannel
+
+        return WechatComChannel()
     raise RuntimeError
diff --git a/channel/wechat/wechat_channel.py b/channel/wechat/wechat_channel.py
index cf200b1..1a1e5ea 100644
--- a/channel/wechat/wechat_channel.py
+++ b/channel/wechat/wechat_channel.py
@@ -29,7 +29,7 @@ from plugins import *
 @itchat.msg_register([TEXT, VOICE, PICTURE, NOTE])
 def handler_single_msg(msg):
     try:
-        cmsg = WeChatMessage(msg, False)
+        cmsg = WechatMessage(msg, False)
     except NotImplementedError as e:
         logger.debug("[WX]single message {} skipped: {}".format(msg["MsgId"], e))
         return None
@@ -40,7 +40,7 @@ def handler_single_msg(msg):
 @itchat.msg_register([TEXT, VOICE, PICTURE, NOTE], isGroupChat=True)
 def handler_group_msg(msg):
     try:
-        cmsg = WeChatMessage(msg, True)
+        cmsg = WechatMessage(msg, True)
     except NotImplementedError as e:
         logger.debug("[WX]group message {} skipped: {}".format(msg["MsgId"], e))
         return None
diff --git a/channel/wechat/wechat_message.py b/channel/wechat/wechat_message.py
index 1888425..526b24f 100644
--- a/channel/wechat/wechat_message.py
+++ b/channel/wechat/wechat_message.py
@@ -8,7 +8,7 @@ from lib import itchat
 from lib.itchat.content import *
 
 
-class WeChatMessage(ChatMessage):
+class WechatMessage(ChatMessage):
     def __init__(self, itchat_msg, is_group=False):
         super().__init__(itchat_msg)
         self.msg_id = itchat_msg["MsgId"]
diff --git a/channel/wechatcom/wechatcom_channel.py b/channel/wechatcom/wechatcom_channel.py
new file mode 100644
index 0000000..1fbeccd
--- /dev/null
+++ b/channel/wechatcom/wechatcom_channel.py
@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+# -*- coding=utf-8 -*-
+import web
+from wechatpy.enterprise import WeChatClient, create_reply, parse_message
+from wechatpy.enterprise.crypto import WeChatCrypto
+from wechatpy.enterprise.exceptions import InvalidCorpIdException
+from wechatpy.exceptions import InvalidSignatureException
+
+from bridge.context import Context
+from bridge.reply import Reply, ReplyType
+from channel.chat_channel import ChatChannel
+from channel.wechatcom.wechatcom_message import WechatComMessage
+from common.log import logger
+from common.singleton import singleton
+from config import conf
+
+
+@singleton
+class WechatComChannel(ChatChannel):
+    NOT_SUPPORT_REPLYTYPE = [ReplyType.IMAGE, ReplyType.VOICE]
+
+    def __init__(self):
+        super().__init__()
+        self.corp_id = conf().get("wechatcom_corp_id")
+        self.secret = conf().get("wechatcom_secret")
+        self.agent_id = conf().get("wechatcom_agent_id")
+        self.token = conf().get("wechatcom_token")
+        self.aes_key = conf().get("wechatcom_aes_key")
+        print(self.corp_id, self.secret, self.agent_id, self.token, self.aes_key)
+        logger.info(
+            "[wechatcom] init: corp_id: {}, secret: {}, agent_id: {}, token: {}, aes_key: {}".format(
+                self.corp_id, self.secret, self.agent_id, self.token, self.aes_key
+            )
+        )
+        self.crypto = WeChatCrypto(self.token, self.aes_key, self.corp_id)
+        self.client = WeChatClient(self.corp_id, self.secret)  # todo: 这里可能有线程安全问题
+
+    def startup(self):
+        # start message listener
+        urls = ("/wxcom", "channel.wechatcom.wechatcom_channel.Query")
+        app = web.application(urls, globals(), autoreload=False)
+        port = conf().get("wechatcom_port", 8080)
+        web.httpserver.runsimple(app.wsgifunc(), ("0.0.0.0", port))
+
+    def send(self, reply: Reply, context: Context):
+        print("send reply: ", reply.content, context["receiver"])
+        receiver = context["receiver"]
+        reply_text = reply.content
+        self.client.message.send_text(self.agent_id, receiver, reply_text)
+        logger.info("[send] Do send to {}: {}".format(receiver, reply_text))
+
+
+class Query:
+    def GET(self):
+        channel = WechatComChannel()
+        params = web.input()
+        signature = params.msg_signature
+        timestamp = params.timestamp
+        nonce = params.nonce
+        echostr = params.echostr
+        print(params)
+        try:
+            echostr = channel.crypto.check_signature(
+                signature, timestamp, nonce, echostr
+            )
+        except InvalidSignatureException:
+            raise web.Forbidden()
+        return echostr
+
+    def POST(self):
+        channel = WechatComChannel()
+        params = web.input()
+        signature = params.msg_signature
+        timestamp = params.timestamp
+        nonce = params.nonce
+        try:
+            message = channel.crypto.decrypt_message(
+                web.data(), signature, timestamp, nonce
+            )
+        except (InvalidSignatureException, InvalidCorpIdException):
+            raise web.Forbidden()
+        print(message)
+        msg = parse_message(message)
+
+        print(msg)
+        if msg.type == "event":
+            if msg.event == "subscribe":
+                reply = create_reply("感谢关注", msg).render()
+                res = channel.crypto.encrypt_message(reply, nonce, timestamp)
+                return res
+        else:
+            try:
+                wechatcom_msg = WechatComMessage(msg, client=channel.client)
+            except NotImplementedError as e:
+                logger.debug("[wechatcom] " + str(e))
+                return "success"
+            context = channel._compose_context(
+                wechatcom_msg.ctype,
+                wechatcom_msg.content,
+                isgroup=False,
+                msg=wechatcom_msg,
+            )
+            if context:
+                channel.produce(context)
+        return "success"
diff --git a/channel/wechatcom/wechatcom_message.py b/channel/wechatcom/wechatcom_message.py
new file mode 100644
index 0000000..42de66e
--- /dev/null
+++ b/channel/wechatcom/wechatcom_message.py
@@ -0,0 +1,54 @@
+import re
+
+import requests
+from wechatpy.enterprise import WeChatClient
+
+from bridge.context import ContextType
+from channel.chat_message import ChatMessage
+from common.log import logger
+from common.tmp_dir import TmpDir
+from lib import itchat
+from lib.itchat.content import *
+
+
+class WechatComMessage(ChatMessage):
+    def __init__(self, msg, client: WeChatClient, is_group=False):
+        super().__init__(msg)
+        self.msg_id = msg.id
+        self.create_time = msg.time
+        self.is_group = is_group
+
+        if msg.type == "text":
+            self.ctype = ContextType.TEXT
+            self.content = msg.content
+        elif msg.type == "voice":
+            self.ctype = ContextType.VOICE
+            self.content = (
+                TmpDir().path() + msg.media_id + "." + msg.format
+            )  # content直接存临时目录路径
+
+            def download_voice():
+                # 如果响应状态码是200,则将响应内容写入本地文件
+                response = client.media.download(msg.media_id)
+                if response.status_code == 200:
+                    with open(self.content, "wb") as f:
+                        f.write(response.content)
+                else:
+                    logger.info(
+                        f"[wechatcom] Failed to download voice file, {response.content}"
+                    )
+
+            self._prepare_fn = download_voice
+        elif msg.type == "image":
+            self.ctype = ContextType.IMAGE
+            self.content = msg.image  # content直接存临时目录路径
+            print(self.content)
+            # self._prepare_fn = lambda: itchat_msg.download(self.content) # TODO: download image
+        else:
+            raise NotImplementedError(
+                "Unsupported message type: Type:{} ".format(msg.type)
+            )
+
+        self.from_user_id = msg.source
+        self.to_user_id = msg.target
+        self.other_user_id = msg.source
diff --git a/channel/wechatmp/wechatmp_channel.py b/channel/wechatmp/wechatmp_channel.py
index ac3c3ac..3152124 100644
--- a/channel/wechatmp/wechatmp_channel.py
+++ b/channel/wechatmp/wechatmp_channel.py
@@ -98,7 +98,9 @@ class WechatMPChannel(ChatChannel):
         if self.passive_reply:
             receiver = context["receiver"]
             self.cache_dict[receiver] = reply.content
-            logger.info("[send] reply to {} saved to cache: {}".format(receiver, reply))
+            logger.info(
+                "[wechatmp] reply to {} saved to cache: {}".format(receiver, reply)
+            )
         else:
             receiver = context["receiver"]
             reply_text = reply.content
@@ -115,7 +117,7 @@ class WechatMPChannel(ChatChannel):
                 params=params,
                 data=json.dumps(json_data, ensure_ascii=False).encode("utf8"),
             )
-            logger.info("[send] Do send to {}: {}".format(receiver, reply_text))
+            logger.info("[wechatmp] Do send to {}: {}".format(receiver, reply_text))
         return
 
     def _success_callback(self, session_id, context, **kwargs):  # 线程异常结束时的回调函数
diff --git a/config.py b/config.py
index 8f5d2ca..f6c3983 100644
--- a/config.py
+++ b/config.py
@@ -75,6 +75,13 @@ available_setting = {
     "wechatmp_port": 8080,  # 微信公众平台的端口,需要端口转发到80或443
     "wechatmp_app_id": "",  # 微信公众平台的appID,仅服务号需要
     "wechatmp_app_secret": "",  # 微信公众平台的appsecret,仅服务号需要
+    # wechatcom的配置
+    "wechatcom_token": "",  # 企业微信的token
+    "wechatcom_port": 9898,  # 企业微信的服务端口,不需要端口转发
+    "wechatcom_corp_id": "",  # 企业微信的corpID
+    "wechatcom_secret": "",  # 企业微信的secret
+    "wechatcom_agent_id": "",  # 企业微信的appID
+    "wechatcom_aes_key": "",  # 企业微信的aes_key
     # chatgpt指令自定义触发词
     "clear_memory_commands": ["#清除记忆"],  # 重置会话指令,必须以#开头
     # channel配置
diff --git a/requirements-optional.txt b/requirements-optional.txt
index cfb52c9..d7c48ac 100644
--- a/requirements-optional.txt
+++ b/requirements-optional.txt
@@ -16,8 +16,9 @@ wechaty>=0.10.7
 wechaty_puppet>=0.4.23
 pysilk_mod>=1.6.0 # needed by send voice
 
-# wechatmp
+# wechatmp wechatcom
 web.py
+wechatpy
 
 # chatgpt-tool-hub plugin