From 6502a64cab043fa84bd1f5acabde7f671850c2b5 Mon Sep 17 00:00:00 2001 From: Junyan Qin Date: Wed, 12 Feb 2025 21:12:53 +0800 Subject: [PATCH] feat(lark): supports for encrypted message --- pkg/core/migrations/m021_lark_config.py | 4 +- pkg/core/migrations/m030_lark_config_cmpl.py | 31 +++++++ pkg/core/stages/migrate.py | 2 +- pkg/platform/sources/lark.py | 94 ++++++++++++++------ templates/platform.json | 4 +- templates/schema/platform.json | 13 ++- 6 files changed, 115 insertions(+), 33 deletions(-) create mode 100644 pkg/core/migrations/m030_lark_config_cmpl.py diff --git a/pkg/core/migrations/m021_lark_config.py b/pkg/core/migrations/m021_lark_config.py index 09adbc42..03d3c9a5 100644 --- a/pkg/core/migrations/m021_lark_config.py +++ b/pkg/core/migrations/m021_lark_config.py @@ -24,7 +24,9 @@ class LarkConfigMigration(migration.Migration): "app_id": "cli_abcdefgh", "app_secret": "XXXXXXXXXX", "bot_name": "LangBot", - "port": None + "enable-webhook": False, + "port": 2285, + "encrypt-key": "xxxxxxxxx" }) await self.ap.platform_cfg.dump_config() diff --git a/pkg/core/migrations/m030_lark_config_cmpl.py b/pkg/core/migrations/m030_lark_config_cmpl.py new file mode 100644 index 00000000..e016af7b --- /dev/null +++ b/pkg/core/migrations/m030_lark_config_cmpl.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from .. import migration + + +@migration.migration_class("lark-config-cmpl", 30) +class LarkConfigCmplMigration(migration.Migration): + """迁移""" + + async def need_migrate(self) -> bool: + """判断当前环境是否需要运行此迁移""" + + for adapter in self.ap.platform_cfg.data['platform-adapters']: + if adapter['adapter'] == 'lark': + if 'enable-webhook' not in adapter: + return True + + return False + + async def run(self): + """执行迁移""" + for adapter in self.ap.platform_cfg.data['platform-adapters']: + if adapter['adapter'] == 'lark': + if 'enable-webhook' not in adapter: + adapter['enable-webhook'] = False + if 'port' not in adapter: + adapter['port'] = 2285 + if 'encrypt-key' not in adapter: + adapter['encrypt-key'] = "xxxxxxxxx" + + await self.ap.platform_cfg.dump_config() diff --git a/pkg/core/stages/migrate.py b/pkg/core/stages/migrate.py index a1983f0b..e144e83a 100644 --- a/pkg/core/stages/migrate.py +++ b/pkg/core/stages/migrate.py @@ -10,7 +10,7 @@ from ..migrations import m010_ollama_requester_config, m011_command_prefix_confi from ..migrations import m015_gitee_ai_config, m016_dify_service_api, m017_dify_api_timeout_params, m018_xai_config, m019_zhipuai_config from ..migrations import m020_wecom_config, m021_lark_config, m022_lmstudio_config, m023_siliconflow_config, m024_discord_config, m025_gewechat_config from ..migrations import m026_qqofficial_config, m027_wx_official_account_config - +from ..migrations import m030_lark_config_cmpl @stage.stage_class("MigrationStage") class MigrationStage(stage.BootingStage): """迁移阶段 diff --git a/pkg/platform/sources/lark.py b/pkg/platform/sources/lark.py index eecf9342..81ae76c5 100644 --- a/pkg/platform/sources/lark.py +++ b/pkg/platform/sources/lark.py @@ -11,6 +11,9 @@ import base64 import uuid import json import datetime +import hashlib +import base64 +from Crypto.Cipher import AES import aiohttp import lark_oapi.ws.exception @@ -28,6 +31,28 @@ from ..types import entities as platform_entities from ...utils import image +class AESCipher(object): + def __init__(self, key): + self.bs = AES.block_size + self.key=hashlib.sha256(AESCipher.str_to_bytes(key)).digest() + @staticmethod + def str_to_bytes(data): + u_type = type(b"".decode('utf8')) + if isinstance(data, u_type): + return data.encode('utf8') + return data + @staticmethod + def _unpad(s): + return s[:-ord(s[len(s) - 1:])] + def decrypt(self, enc): + iv = enc[:AES.block_size] + cipher = AES.new(self.key, AES.MODE_CBC, iv) + return self._unpad(cipher.decrypt(enc[AES.block_size:])) + def decrypt_string(self, enc): + enc = base64.b64decode(enc) + return self.decrypt(enc).decode('utf8') + + class LarkMessageConverter(adapter.MessageConverter): @staticmethod @@ -301,37 +326,47 @@ class LarkMessageSourceAdapter(adapter.MessageSourceAdapter): @self.quart_app.route('/lark/callback', methods=['POST']) async def lark_callback(): - data = await quart.request.json + try: + data = await quart.request.json - type = data.get("type") - if type is None : + if 'encrypt' in data: + cipher = AESCipher(self.config['encrypt-key']) + data = cipher.decrypt_string(data['encrypt']) + data = json.loads(data) + + type = data.get("type") + if type is None : + context = EventContext(data) + type = context.header.event_type + + if 'url_verification' == type: + print(data.get("challenge")) + # todo 验证verification token + return { + "challenge": data.get("challenge") + } context = EventContext(data) type = context.header.event_type - if type is None : - return 'ok' + p2v1 = P2ImMessageReceiveV1() + p2v1.header = context.header + event = P2ImMessageReceiveV1Data() + event.message = EventMessage(context.event['message']) + event.sender = EventSender(context.event['sender']) + p2v1.event = event + p2v1.schema = context.schema + if 'im.message.receive_v1' == type: + try: + event = await self.event_converter.target2yiri(p2v1, self.api_client) + except Exception as e: + traceback.print_exc() - if 'url_verification' in type: - # todo 验证verification token - return data - context = EventContext(data) - type = context.header.event_type - p2v1 = P2ImMessageReceiveV1() - p2v1.header = context.header - event = P2ImMessageReceiveV1Data() - event.message = EventMessage(context.event['message']) - event.sender = EventSender(context.event['sender']) - p2v1.event = event - p2v1.schema = context.schema - if 'im.message.receive_v1' in type: - try: - event = await self.event_converter.target2yiri(p2v1, self.api_client) - except Exception as e: - traceback.print_exc() + if event.__class__ in self.listeners: + await self.listeners[event.__class__](event, self) - if event.__class__ in self.listeners: - await self.listeners[event.__class__](event, self) - - return 'ok' + return {"code": 200, "message": "ok"} + except Exception as e: + traceback.print_exc() + return {"code": 500, "message": "error"} async def on_message(event: lark_oapi.im.v1.P2ImMessageReceiveV1): @@ -430,9 +465,10 @@ class LarkMessageSourceAdapter(adapter.MessageSourceAdapter): self.listeners.pop(event_type) async def run_async(self): - port =self.config['port'] + port = self.config['port'] + enable_webhook = self.config['enable-webhook'] - if port is None or port == "": + if not enable_webhook: try: await self.bot._connect() except lark_oapi.ws.exception.ClientException as e: @@ -450,7 +486,7 @@ class LarkMessageSourceAdapter(adapter.MessageSourceAdapter): await self.quart_app.run_task( host='0.0.0.0', - port=self.config['port'], + port=port, shutdown_trigger=shutdown_trigger_placeholder, ) async def kill(self) -> bool: diff --git a/templates/platform.json b/templates/platform.json index ed2bf571..703a2ad4 100644 --- a/templates/platform.json +++ b/templates/platform.json @@ -50,7 +50,9 @@ "app_id": "cli_abcdefgh", "app_secret": "XXXXXXXXXX", "bot_name": "LangBot", - "port":"" + "enable-webhook": false, + "port": 2285, + "encrypt-key": "xxxxxxxxx" }, { "adapter": "discord", diff --git a/templates/schema/platform.json b/templates/schema/platform.json index bc5d9962..f0224e61 100644 --- a/templates/schema/platform.json +++ b/templates/schema/platform.json @@ -253,9 +253,20 @@ "default": "", "description": "飞书的bot_name" }, + "enable-webhook": { + "type": "boolean", + "default": false, + "description": "是否启用webhook模式" + }, "port": { "type": "integer", - "description": "设置监听的端口,开启callback event" + "description": "设置监听的端口,开启callback event时需要设置", + "default": 2285 + }, + "encrypt-key": { + "type": "string", + "default": "", + "description": "设置加密密钥" } } },