246 lines
9.0 KiB
Python
246 lines
9.0 KiB
Python
"""钉钉 IM 适配器 (U12)。
|
||
|
||
实现 :class:`MessageAdapter` 协议，对接钉钉企业内机器人 outgoing/webhook 回调。
|
||
|
||
关键设计决策：
|
||
- 签名校验 fail-closed：``Sign`` + ``Timestamp`` 头缺失时仅当 token 校验通过才放行。
|
||
- 时间戳窗口 3600s（钉钉官方窗口 1 小时）。
|
||
- ``accessToken`` 简单 TTL 缓存（7200s）；钉钉 token 有效期 2 小时。
|
||
- httpx 客户端懒构造，避免未使用的适配器持有连接池。
|
||
- 钉钉无 URL verification challenge 流程 — 合法签名请求直接 200。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import hmac
|
||
import json
|
||
import logging
|
||
import re
|
||
import time
|
||
|
||
import httpx
|
||
|
||
from agentkit.channels.base import (
|
||
ChannelType,
|
||
IncomingMessage,
|
||
MessageAdapter,
|
||
OutgoingMessage,
|
||
header_get,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 签名时间戳允许的最大偏移（秒）— 钉钉官方窗口 1 小时
|
||
_SIGNATURE_MAX_AGE_SECONDS = 3600
|
||
# accessToken 缓存 TTL（秒）— 钉钉 token 有效期 2 小时
|
||
_TOKEN_CACHE_TTL = 7200.0
|
||
|
||
# 钉钉 API 端点
|
||
_ACCESS_TOKEN_URL = "https://api.dingtalk.com/v1.0/oauth2/accessToken"
|
||
_SEND_MESSAGE_URL = "https://api.dingtalk.com/v1.0/robot/oToMessages/batchSend"
|
||
|
||
# 钉钉 @ 机器人前缀剥离 — 内容形如 "@robotName 实际消息"
|
||
# ponytail: 简单正则剥离首个 @token。天花板：无法区分 @ 机器人与 @ 用户；
|
||
# 升级路径：从 config 注入 robot_name 精确匹配。
|
||
_MENTION_PREFIX_RE = re.compile(r"^@\S+\s*")
|
||
|
||
|
||
class DingTalkMessageAdapter(MessageAdapter):
|
||
"""钉钉 IM 适配器。
|
||
|
||
生命周期：
|
||
``__init__`` → :meth:`verify_signature` → :meth:`receive_message`
|
||
→ :meth:`send_message` → :meth:`close`
|
||
|
||
Args:
|
||
app_key: 钉钉应用 App Key。
|
||
app_secret: 钉钉应用 App Secret（同时作为签名密钥）。
|
||
robot_code: 机器人 robotCode（发送消息时使用）。
|
||
token: 可选 token 校验值（部分机器人通过 ``Token`` 头校验）。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
app_key: str,
|
||
app_secret: str,
|
||
robot_code: str,
|
||
token: str | None = None,
|
||
) -> None:
|
||
super().__init__()
|
||
self.app_key = app_key
|
||
self.app_secret = app_secret
|
||
self.robot_code = robot_code
|
||
self.token = token
|
||
# ponytail: 简单 TTL 缓存 (token, expiry)。天花板：单实例内存；
|
||
# 升级路径：Redis 缓存共享给多实例。
|
||
self._token_cache: tuple[str, float] | None = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# 签名验证
|
||
# ------------------------------------------------------------------
|
||
|
||
async def verify_signature(self, headers: dict[str, str], body: bytes) -> bool:
|
||
"""验证钉钉 webhook 签名。
|
||
|
||
- 若配置了 ``token``：校验 ``Token`` 头，不匹配返回 False。
|
||
- 若存在 ``Sign`` + ``Timestamp`` 头：校验 HMAC-SHA256 签名与时间戳新鲜度。
|
||
- 两者皆无：仅当 token 已校验通过才放行（token-only 模式）。
|
||
|
||
Args:
|
||
headers: HTTP 请求头（键大小写不敏感查找）。
|
||
body: 原始请求体字节（钉钉签名不依赖 body）。
|
||
|
||
Returns:
|
||
True 表示签名校验通过。
|
||
"""
|
||
# Token 校验（若配置）
|
||
if self.token is not None:
|
||
token_header = header_get(headers, "Token")
|
||
if token_header is None or not hmac.compare_digest(token_header, self.token):
|
||
return False
|
||
|
||
sign = header_get(headers, "Sign")
|
||
timestamp_str = header_get(headers, "Timestamp")
|
||
|
||
if sign is None and timestamp_str is None:
|
||
# 无签名头：仅当 token 已校验通过才放行
|
||
return self.token is not None
|
||
|
||
if sign is None or timestamp_str is None:
|
||
return False
|
||
|
||
# 时间戳新鲜度（毫秒 → 秒）
|
||
try:
|
||
ts_ms = int(timestamp_str)
|
||
except ValueError:
|
||
return False
|
||
ts_sec = ts_ms / 1000.0
|
||
now = time.time()
|
||
if abs(now - ts_sec) > _SIGNATURE_MAX_AGE_SECONDS:
|
||
logger.warning("钉钉 webhook 时间戳超出 %ds 窗口 — 拒绝", _SIGNATURE_MAX_AGE_SECONDS)
|
||
return False
|
||
|
||
# 计算签名：base64(hmac-sha256(key=app_secret, msg="{timestamp}\n{app_secret}"))
|
||
string_to_sign = f"{timestamp_str}\n{self.app_secret}"
|
||
expected = base64.b64encode(
|
||
hmac.new(
|
||
self.app_secret.encode("utf-8"),
|
||
string_to_sign.encode("utf-8"),
|
||
hashlib.sha256,
|
||
).digest()
|
||
).decode("utf-8")
|
||
return hmac.compare_digest(sign, expected)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 消息接收 / 解析
|
||
# ------------------------------------------------------------------
|
||
|
||
async def receive_message(self, headers: dict[str, str], body: bytes) -> IncomingMessage:
|
||
"""解析钉钉 webhook 事件为标准化 :class:`IncomingMessage`。
|
||
|
||
钉钉无 URL verification challenge 流程 — 合法请求直接解析为消息。
|
||
|
||
Raises:
|
||
ValueError: 事件 body 不是合法 JSON。
|
||
"""
|
||
try:
|
||
data: dict[str, object] = json.loads(body)
|
||
except json.JSONDecodeError as exc:
|
||
raise ValueError(f"钉钉事件 body 不是合法 JSON: {exc}") from exc
|
||
|
||
chat_id = data.get("conversationId", "")
|
||
user_id = data.get("senderStaffId") or data.get("senderId") or data.get("staffId") or ""
|
||
msg_id = data.get("msgId", "")
|
||
session_expired = data.get("sessionWebhookExpiredTime", "")
|
||
timestamp = str(session_expired) if session_expired else ""
|
||
|
||
content = self._extract_content(data)
|
||
|
||
return IncomingMessage(
|
||
channel=ChannelType.DINGTALK,
|
||
platform_message_id=msg_id,
|
||
user_id=user_id,
|
||
chat_id=chat_id,
|
||
content=content,
|
||
raw_event=data,
|
||
timestamp=timestamp,
|
||
)
|
||
|
||
def _extract_content(self, data: dict[str, object]) -> str:
|
||
"""从钉钉事件提取文本内容。
|
||
|
||
- text 类型：解析 ``text.content``，剥离 @ 机器人前缀。
|
||
- 其他类型：返回 ``[unsupported message type: {type}]``。
|
||
"""
|
||
msgtype = data.get("msgtype", "")
|
||
if msgtype == "text":
|
||
text_obj = data.get("text", {})
|
||
content = text_obj.get("content", "") if isinstance(text_obj, dict) else ""
|
||
# 剥离 @ 机器人前缀
|
||
return _MENTION_PREFIX_RE.sub("", content, count=1).strip()
|
||
return f"[unsupported message type: {msgtype}]"
|
||
|
||
# ------------------------------------------------------------------
|
||
# 消息发送
|
||
# ------------------------------------------------------------------
|
||
|
||
async def send_message(self, message: OutgoingMessage) -> bool:
|
||
"""向钉钉发送文本消息（oToMessages/batchSend 单聊）。
|
||
|
||
Returns:
|
||
True 表示 HTTP 200。
|
||
"""
|
||
try:
|
||
token = await self._get_access_token()
|
||
if not token:
|
||
return False
|
||
|
||
client = self._get_client()
|
||
payload = {
|
||
"robotCode": self.robot_code,
|
||
"conversationId": message.chat_id,
|
||
"msgKey": "sampleText",
|
||
"msgParam": json.dumps({"content": message.content}),
|
||
}
|
||
resp = await client.post(
|
||
_SEND_MESSAGE_URL,
|
||
json=payload,
|
||
headers={"x-acs-dingtalk-access-token": token},
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.error("钉钉 send_message HTTP %d: %s", resp.status_code, resp.text[:200])
|
||
return False
|
||
return True
|
||
except httpx.HTTPError as exc:
|
||
logger.error("钉钉 send_message 网络错误: %s", exc)
|
||
return False
|
||
|
||
async def _get_access_token(self) -> str | None:
|
||
"""获取并缓存钉钉 ``accessToken``。"""
|
||
# 命中缓存
|
||
if self._token_cache is not None:
|
||
token, expiry = self._token_cache
|
||
if time.monotonic() < expiry:
|
||
return token
|
||
|
||
try:
|
||
client = self._get_client()
|
||
resp = await client.post(
|
||
_ACCESS_TOKEN_URL,
|
||
json={"appKey": self.app_key, "appSecret": self.app_secret},
|
||
)
|
||
if resp.status_code != 200:
|
||
logger.error("钉钉 accessToken HTTP %d: %s", resp.status_code, resp.text[:200])
|
||
return None
|
||
data = resp.json()
|
||
token = data.get("accessToken", "")
|
||
if not token:
|
||
return None
|
||
self._token_cache = (token, time.monotonic() + _TOKEN_CACHE_TTL)
|
||
return token
|
||
except httpx.HTTPError as exc:
|
||
logger.error("钉钉 accessToken 网络错误: %s", exc)
|
||
return None
|