509 lines
21 KiB
Python
509 lines
21 KiB
Python
"""企微服务：消息加解密、异步处理、API 调用。"""
|
||
import hashlib
|
||
import xml.etree.ElementTree as ET
|
||
from threading import Thread, Lock
|
||
import base64
|
||
import struct
|
||
import time
|
||
import requests
|
||
from insurance.config import get_config as _get_wecom_config
|
||
|
||
|
||
class WeComService:
|
||
"""企微机器人业务逻辑。"""
|
||
|
||
# access_token缓存（加锁防竞态）
|
||
_access_token = None
|
||
_access_token_expire = 0
|
||
_token_lock = Lock()
|
||
|
||
# 消息排重缓存：{dedup_key: expire_timestamp}
|
||
_processed_msgs = {}
|
||
_DEDUP_TTL = 300 # 5分钟过期（企微重试3次，间隔较短）
|
||
|
||
# 回调IP白名单缓存
|
||
_callback_ips = None
|
||
_callback_ips_expire = 0
|
||
_IP_CACHE_TTL = 3600 # 1小时刷新一次
|
||
|
||
def _get_access_token(self) -> str:
|
||
"""获取企微access_token（带缓存+锁）。"""
|
||
# 快速路径：缓存有效直接返回
|
||
if self._access_token and time.time() < self._access_token_expire:
|
||
return self._access_token
|
||
|
||
with self._token_lock:
|
||
# 双重检查：等锁期间可能已被其他线程刷新
|
||
if self._access_token and time.time() < self._access_token_expire:
|
||
return self._access_token
|
||
|
||
corp_id = _get_wecom_config("WECOM_CORP_ID")
|
||
secret = _get_wecom_config("WECOM_SECRET")
|
||
|
||
token_resp = requests.get(
|
||
"https://qyapi.weixin.qq.com/cgi-bin/gettoken",
|
||
params={"corpid": corp_id, "corpsecret": secret},
|
||
timeout=5,
|
||
)
|
||
|
||
resp_data = token_resp.json()
|
||
access_token = resp_data.get("access_token", "")
|
||
expires_in = resp_data.get("expires_in", 7200)
|
||
|
||
if not access_token:
|
||
print(f"[WECOM ERROR] 获取access_token失败: {resp_data}", flush=True)
|
||
return ""
|
||
|
||
# 缓存access_token（提前5分钟过期）
|
||
self._access_token = access_token
|
||
self._access_token_expire = time.time() + expires_in - 300
|
||
|
||
print(f"[WECOM DEBUG] 获取access_token成功，有效期: {expires_in}秒", flush=True)
|
||
return access_token
|
||
|
||
def verify_url(self, msg_signature: str, timestamp: str, nonce: str, echostr: str) -> str:
|
||
"""企微服务器 URL 验证：SHA1 校验 + 解密 echostr。"""
|
||
token = _get_wecom_config("WECOM_TOKEN")
|
||
aes_key = _get_wecom_config("WECOM_AES_KEY")
|
||
corp_id = _get_wecom_config("WECOM_CORP_ID")
|
||
|
||
# 使用print确保输出到stdout
|
||
print(f"[WECOM DEBUG] msg_signature={msg_signature}", flush=True)
|
||
print(f"[WECOM DEBUG] timestamp={timestamp}, nonce={nonce}", flush=True)
|
||
print(f"[WECOM DEBUG] Token={token}, AES_KEY={aes_key[:10]}..., CorpID={corp_id}", flush=True)
|
||
|
||
# 1. 字典序排序拼接
|
||
sort_list = sorted([token, timestamp, nonce, echostr])
|
||
sha1 = hashlib.sha1("".join(sort_list).encode()).hexdigest()
|
||
print(f"[WECOM DEBUG] 计算签名: {sha1}", flush=True)
|
||
print(f"[WECOM DEBUG] 收到签名: {msg_signature}", flush=True)
|
||
|
||
# 2. 校验签名
|
||
if sha1 != msg_signature:
|
||
print(f"[WECOM ERROR] 签名验证失败!", flush=True)
|
||
return "Invalid signature"
|
||
|
||
print(f"[WECOM DEBUG] 签名验证成功", flush=True)
|
||
|
||
# 3. 解密 echostr
|
||
try:
|
||
decrypted = self._aes_decrypt(aes_key, echostr)
|
||
print(f"[WECOM DEBUG] 解密成功: {decrypted[:100]}", flush=True)
|
||
|
||
# 企微URL验证返回的是明文，直接返回
|
||
print(f"[WECOM DEBUG] 返回echostr: {decrypted}", flush=True)
|
||
return decrypted
|
||
except Exception as e:
|
||
print(f"[WECOM ERROR] 解密失败: {e}", flush=True)
|
||
return "Invalid echostr"
|
||
|
||
def _is_duplicate(self, dedup_key: str) -> bool:
|
||
"""检查消息是否已处理过（排重）。"""
|
||
now = time.time()
|
||
# 清理过期条目
|
||
expired = [k for k, v in self._processed_msgs.items() if v < now]
|
||
for k in expired:
|
||
del self._processed_msgs[k]
|
||
# 检查是否重复
|
||
if dedup_key in self._processed_msgs:
|
||
return True
|
||
self._processed_msgs[dedup_key] = now + self._DEDUP_TTL
|
||
return False
|
||
|
||
def _get_callback_ips(self) -> list:
|
||
"""获取企微回调IP白名单（带缓存）。"""
|
||
now = time.time()
|
||
if self._callback_ips is not None and now < self._callback_ips_expire:
|
||
return self._callback_ips
|
||
|
||
access_token = self._get_access_token()
|
||
if not access_token:
|
||
return self._callback_ips or []
|
||
|
||
try:
|
||
resp = requests.get(
|
||
"https://qyapi.weixin.qq.com/cgi-bin/getcallbackip",
|
||
params={"access_token": access_token},
|
||
timeout=5,
|
||
)
|
||
data = resp.json()
|
||
if data.get("errcode") == 0:
|
||
self._callback_ips = data.get("ip_list", [])
|
||
self._callback_ips_expire = now + self._IP_CACHE_TTL
|
||
print(f"[WECOM DEBUG] 获取回调IP白名单成功: {len(self._callback_ips)}条", flush=True)
|
||
except Exception as e:
|
||
print(f"[WECOM ERROR] 获取回调IP白名单失败: {e}", flush=True)
|
||
|
||
return self._callback_ips or []
|
||
|
||
def is_callback_ip(self, client_ip: str) -> bool:
|
||
"""检查请求IP是否在企微回调白名单中。"""
|
||
ip_list = self._get_callback_ips()
|
||
if not ip_list:
|
||
# 白名单为空（获取失败），放行以免阻断所有请求
|
||
return True
|
||
|
||
for pattern in ip_list:
|
||
# 企微返回格式如 "101.226.103.*"，按前缀匹配
|
||
prefix = pattern.rstrip("*").rstrip(".")
|
||
if client_ip.startswith(prefix):
|
||
return True
|
||
|
||
return False
|
||
|
||
def handle_message_for_group(self, msg_signature: str, timestamp: str, nonce: str, encrypted_data: bytes) -> bool:
|
||
"""判断是否为群聊消息，若是则启动异步处理线程。
|
||
|
||
Returns:
|
||
True 表示是群聊（路由应返回 "success" 并异步处理）。
|
||
"""
|
||
try:
|
||
decrypted_str = self._decrypt_message(encrypted_data, msg_signature, timestamp, nonce)
|
||
msg_data = self._parse_message(decrypted_str)
|
||
|
||
if not msg_data or msg_data.get("chattype") != "group":
|
||
return False
|
||
|
||
# 群聊：启动异步处理线程
|
||
print(f"[WECOM DEBUG] 检测到群聊，启动异步处理线程", flush=True)
|
||
from flask import copy_current_request_context
|
||
@copy_current_request_context
|
||
def _process_group():
|
||
self._handle_group_message(msg_data)
|
||
t = Thread(target=_process_group, daemon=True)
|
||
t.start()
|
||
print(f"[WECOM DEBUG] 异步处理线程已启动, is_alive={t.is_alive()}", flush=True)
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"[WECOM ERROR] 群聊检测失败: {e}", flush=True)
|
||
return False
|
||
|
||
def handle_message(self, msg_signature: str, timestamp: str, nonce: str, encrypted_data: bytes) -> str:
|
||
"""同步处理私聊消息，返回被动回复 XML。"""
|
||
try:
|
||
decrypted_str = self._decrypt_message(encrypted_data, msg_signature, timestamp, nonce)
|
||
msg_data = self._parse_message(decrypted_str)
|
||
if not msg_data:
|
||
return ""
|
||
reply_text = self._call_baodan(msg_data)
|
||
if reply_text:
|
||
return self._build_passive_reply(reply_text, timestamp, nonce)
|
||
return ""
|
||
except Exception as e:
|
||
print(f"[WECOM ERROR] 消息处理异常: {e}", flush=True)
|
||
return ""
|
||
|
||
def _parse_message(self, decrypted_str: str) -> dict | None:
|
||
"""解析解密后的消息（JSON 或 XML），返回统一格式 dict。"""
|
||
print(f"[WECOM DEBUG] 解密后消息: {decrypted_str[:200]}", flush=True)
|
||
|
||
try:
|
||
import json
|
||
msg_data = json.loads(decrypted_str)
|
||
msg_type = msg_data.get("msgtype", "")
|
||
from_user = msg_data.get("from", {}).get("userid", "")
|
||
chat_type = msg_data.get("chattype", "")
|
||
chat_id = msg_data.get("chatid", "")
|
||
content = msg_data.get("text", {}).get("content", "") if msg_type == "text" else ""
|
||
dedup_key = msg_data.get("msgid", "") or msg_data.get("msg_id", "")
|
||
print(f"[WECOM DEBUG] JSON格式: type={msg_type}, user={from_user}, chat_type={chat_type}, content={content[:50]}", flush=True)
|
||
except (json.JSONDecodeError, AttributeError):
|
||
root = ET.fromstring(decrypted_str)
|
||
msg_type = root.findtext("MsgType", "")
|
||
from_user = root.findtext("FromUserName", "")
|
||
content = root.findtext("Content", "")
|
||
chat_type = ""
|
||
chat_id = ""
|
||
dedup_key = root.findtext("MsgId", "")
|
||
print(f"[WECOM DEBUG] XML格式: type={msg_type}, user={from_user}, content={content[:50]}", flush=True)
|
||
|
||
return {
|
||
"msgtype": msg_type,
|
||
"from_user": from_user,
|
||
"chattype": chat_type,
|
||
"chatid": chat_id,
|
||
"content": content,
|
||
"dedup_key": dedup_key,
|
||
}
|
||
|
||
def _call_baodan(self, msg_data: dict) -> str:
|
||
"""调 BaoDan Chat API 获取回答。"""
|
||
msg_type = msg_data["msgtype"]
|
||
from_user = msg_data["from_user"]
|
||
content = msg_data["content"]
|
||
dedup_key = msg_data["dedup_key"]
|
||
|
||
if dedup_key and self._is_duplicate(dedup_key):
|
||
print(f"[WECOM DEBUG] 重复消息，跳过: {dedup_key}", flush=True)
|
||
return ""
|
||
|
||
if msg_type != "text":
|
||
print(f"[WECOM DEBUG] 非文本消息，回复暂不支持", flush=True)
|
||
return "暂不支持该消息类型，请发送文字消息。"
|
||
|
||
# 去掉群聊中"@机器人"前缀
|
||
if content.startswith("@"):
|
||
content = content.split(" ", 1)[-1] if " " in content else content
|
||
|
||
print(f"[WECOM DEBUG] 处理文本消息: {content}", flush=True)
|
||
|
||
api_key = _get_wecom_config("BAODAN_CHAT_API_KEY")
|
||
base_url = _get_wecom_config("BAODAN_API_URL", "http://baodan-api:5001")
|
||
|
||
print(f"[WECOM DEBUG] 调用API: {base_url}/v1/chat-messages", flush=True)
|
||
|
||
resp = requests.post(
|
||
f"{base_url}/v1/chat-messages",
|
||
json={
|
||
"inputs": {},
|
||
"query": content,
|
||
"response_mode": "blocking",
|
||
"user": f"wecom_{from_user}",
|
||
"conversation_id": "",
|
||
},
|
||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||
timeout=30,
|
||
)
|
||
|
||
answer = resp.json().get("answer", "抱歉，暂时无法回答您的问题。")
|
||
print(f"[WECOM DEBUG] 回答: {answer[:100]}", flush=True)
|
||
return answer
|
||
|
||
def _handle_group_message(self, msg_data: dict):
|
||
"""异步处理群聊消息：调 BaoDan -> 主动发送到群。"""
|
||
print(f"[WECOM DEBUG] 群聊异步处理开始: user={msg_data.get('from_user')}, chatid={msg_data.get('chatid')}", flush=True)
|
||
try:
|
||
answer = self._call_baodan(msg_data)
|
||
print(f"[WECOM DEBUG] 群聊异步处理, 拿到回答: {answer[:100] if answer else '空'}", flush=True)
|
||
if answer:
|
||
chat_id = msg_data.get("chatid", "")
|
||
if chat_id:
|
||
self._send_text_to_group(chat_id, answer)
|
||
else:
|
||
print(f"[WECOM ERROR] 群聊异步处理: chatid 为空", flush=True)
|
||
except Exception as e:
|
||
print(f"[WECOM ERROR] 群聊消息处理失败: {e}", flush=True)
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
def _decrypt_message(self, encrypted_data: bytes, msg_signature: str, timestamp: str, nonce: str) -> str:
|
||
"""解密企微消息。"""
|
||
aes_key = _get_wecom_config("WECOM_AES_KEY")
|
||
token = _get_wecom_config("WECOM_TOKEN")
|
||
|
||
# 调试：查看原始数据
|
||
raw_data = encrypted_data.decode("utf-8", errors="replace")
|
||
print(f"[WECOM DEBUG] 原始数据: {raw_data[:200]}", flush=True)
|
||
|
||
# 1. 提取加密的消息体（支持JSON和XML格式）
|
||
encrypt = ""
|
||
try:
|
||
# 尝试JSON格式（企微新版）
|
||
import json
|
||
data = json.loads(raw_data)
|
||
encrypt = data.get("encrypt", "")
|
||
print(f"[WECOM DEBUG] JSON格式，提取到encrypt: {encrypt[:50] if encrypt else 'None'}", flush=True)
|
||
except (json.JSONDecodeError, AttributeError):
|
||
try:
|
||
# 尝试XML格式（企微旧版）
|
||
root = ET.fromstring(raw_data)
|
||
encrypt = root.findtext("Encrypt", "")
|
||
print(f"[WECOM DEBUG] XML格式，提取到Encrypt: {encrypt[:50] if encrypt else 'None'}", flush=True)
|
||
except ET.ParseError as e:
|
||
print(f"[WECOM ERROR] 解析失败: {e}", flush=True)
|
||
# 尝试将整个数据作为加密消息体处理
|
||
encrypt = raw_data
|
||
|
||
if not encrypt:
|
||
raise ValueError("No Encrypt field found")
|
||
|
||
# 2. 验证签名（使用加密的消息体）
|
||
sort_list = sorted([token, timestamp, nonce, encrypt])
|
||
sha1 = hashlib.sha1("".join(sort_list).encode()).hexdigest()
|
||
print(f"[WECOM DEBUG] 签名验证: 计算={sha1}, 收到={msg_signature}", flush=True)
|
||
|
||
if sha1 != msg_signature:
|
||
raise ValueError("Invalid signature")
|
||
|
||
# 3. 解密消息
|
||
decrypted = self._aes_decrypt(aes_key, encrypt)
|
||
print(f"[WECOM DEBUG] 解密后消息: {decrypted[:200]}", flush=True)
|
||
return decrypted
|
||
|
||
def _aes_decrypt(self, key: str, encrypted: str) -> str:
|
||
"""AES-256-CBC 解密企微消息（企微专用格式）。"""
|
||
# 1. Base64 解码密钥（企微密钥是43位，需要补=号）
|
||
aes_key = base64.b64decode(key + "=")
|
||
|
||
# 2. Base64 解码密文
|
||
cipher_text = base64.b64decode(encrypted)
|
||
|
||
# 3. AES-256-CBC 解密（IV = 密钥前16字节）
|
||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||
from cryptography.hazmat.backends import default_backend
|
||
|
||
iv = aes_key[:16]
|
||
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
||
decryptor = cipher.decryptor()
|
||
decrypted = decryptor.update(cipher_text) + decryptor.finalize()
|
||
|
||
# 4. 去除 PKCS7 填充
|
||
pad_len = decrypted[-1]
|
||
if isinstance(pad_len, int) and 1 <= pad_len <= 32:
|
||
decrypted = decrypted[:-pad_len]
|
||
|
||
# 5. 提取消息内容（企微格式：随机16字节 + 消息长度4字节 + 消息内容 + CorpID）
|
||
# 消息长度在第16-20字节（大端序）
|
||
msg_len = struct.unpack("!I", decrypted[16:20])[0]
|
||
msg_content = decrypted[20:20 + msg_len].decode("utf-8")
|
||
|
||
return msg_content
|
||
|
||
def _aes_encrypt(self, key: str, reply_msg: str) -> str:
|
||
"""AES-256-CBC 加密企微被动回复消息。"""
|
||
import os
|
||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||
from cryptography.hazmat.backends import default_backend
|
||
|
||
aes_key = base64.b64decode(key + "=")
|
||
iv = aes_key[:16]
|
||
|
||
# 构造明文：随机16字节 + 消息长度4字节 + 消息内容 + CorpID
|
||
corp_id = _get_wecom_config("WECOM_CORP_ID")
|
||
msg_bytes = reply_msg.encode("utf-8")
|
||
corp_id_bytes = corp_id.encode("utf-8")
|
||
random_bytes = os.urandom(16)
|
||
msg_len = struct.pack("!I", len(msg_bytes))
|
||
plain = random_bytes + msg_len + msg_bytes + corp_id_bytes
|
||
|
||
# PKCS7 填充
|
||
block_size = 32
|
||
pad_len = block_size - (len(plain) % block_size)
|
||
plain += bytes([pad_len] * pad_len)
|
||
|
||
# AES-256-CBC 加密
|
||
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
||
encryptor = cipher.encryptor()
|
||
encrypted = encryptor.update(plain) + encryptor.finalize()
|
||
|
||
return base64.b64encode(encrypted).decode("utf-8")
|
||
|
||
def _build_passive_reply(self, reply_msg: str, timestamp: str, nonce: str) -> str:
|
||
"""构建企微被动回复的加密 XML 响应。"""
|
||
aes_key = _get_wecom_config("WECOM_AES_KEY")
|
||
token = _get_wecom_config("WECOM_TOKEN")
|
||
|
||
# 1. 加密回复消息
|
||
encrypt = self._aes_encrypt(aes_key, reply_msg)
|
||
|
||
# 2. 生成签名
|
||
sort_list = sorted([token, timestamp, nonce, encrypt])
|
||
msg_signature = hashlib.sha1("".join(sort_list).encode()).hexdigest()
|
||
|
||
# 3. 构造 XML 响应（CreateTime 为必填字段，缺少会导致 WeChat 静默丢弃消息）
|
||
create_time = int(time.time())
|
||
reply_xml = f"""<xml>
|
||
<Encrypt><![CDATA[{encrypt}]]></Encrypt>
|
||
<MsgSignature><![CDATA[{msg_signature}]]></MsgSignature>
|
||
<CreateTime>{create_time}</CreateTime>
|
||
<Nonce><![CDATA[{nonce}]]></Nonce>
|
||
</xml>"""
|
||
return reply_xml
|
||
|
||
def _send_text(self, to_user: str, content: str):
|
||
"""通过企微 API 发送文本消息给个人。"""
|
||
import requests
|
||
|
||
access_token = self._get_access_token()
|
||
if not access_token:
|
||
return
|
||
|
||
agent_id = _get_wecom_config("WECOM_AGENT_ID", "1000002")
|
||
|
||
# 企微消息长度限制 2048 字符
|
||
MAX_LEN = 2000
|
||
chunks = [content[i:i + MAX_LEN] for i in range(0, len(content), MAX_LEN)]
|
||
|
||
for chunk in chunks:
|
||
resp = requests.post(
|
||
f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={access_token}",
|
||
json={
|
||
"touser": to_user,
|
||
"msgtype": "text",
|
||
"agentid": int(agent_id),
|
||
"text": {"content": chunk},
|
||
},
|
||
timeout=10,
|
||
)
|
||
resp_data = resp.json()
|
||
if resp_data.get("errcode") != 0:
|
||
print(f"[WECOM ERROR] 发送消息失败: {resp_data}", flush=True)
|
||
else:
|
||
print(f"[WECOM DEBUG] 发送消息给{to_user}成功", flush=True)
|
||
|
||
def _send_text_to_group(self, chat_id: str, content: str):
|
||
"""通过企微应用消息接口发送文本消息到群聊。
|
||
|
||
注意：企微的应用消息接口 /cgi-bin/message/send 发送群聊消息时，
|
||
需要在企微后台开启"群聊会话"功能，并配置接收消息的群聊。
|
||
如果使用群机器人 Webhook，需要配置 WECOM_GROUP_WEBHOOK 环境变量。
|
||
"""
|
||
# 方式1：优先使用群机器人 Webhook（更稳定）
|
||
webhook_url = _get_wecom_config("WECOM_GROUP_WEBHOOK")
|
||
if webhook_url:
|
||
self._send_text_to_group_webhook(webhook_url, content)
|
||
return
|
||
|
||
# 方式2：使用应用消息接口（需要配置"群聊会话"功能）
|
||
access_token = self._get_access_token()
|
||
if not access_token:
|
||
return
|
||
|
||
agent_id = _get_wecom_config("WECOM_AGENT_ID", "1000002")
|
||
print(f"[WECOM DEBUG] 发送群聊消息: chat_id={chat_id}, agent_id={agent_id}", flush=True)
|
||
|
||
# 企微消息长度限制 2048 字符
|
||
MAX_LEN = 2000
|
||
chunks = [content[i:i + MAX_LEN] for i in range(0, len(content), MAX_LEN)]
|
||
|
||
for chunk in chunks:
|
||
# 应用消息推送接口，接收者类型为 chatid
|
||
resp = requests.post(
|
||
f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={access_token}",
|
||
json={
|
||
"chatid": chat_id,
|
||
"msgtype": "text",
|
||
"agentid": int(agent_id),
|
||
"text": {"content": chunk},
|
||
},
|
||
timeout=10,
|
||
)
|
||
resp_data = resp.json()
|
||
print(f"[WECOM DEBUG] 群聊消息发送响应: {resp_data}", flush=True)
|
||
if resp_data.get("errcode") != 0:
|
||
print(f"[WECOM ERROR] 发送消息到群失败: {resp_data}", flush=True)
|
||
else:
|
||
print(f"[WECOM DEBUG] 发送消息到群{chat_id}成功", flush=True)
|
||
|
||
def _send_text_to_group_webhook(self, webhook_url: str, content: str):
|
||
"""通过群机器人 Webhook 发送文本消息到群聊。"""
|
||
# 企微消息长度限制 2048 字符
|
||
MAX_LEN = 2000
|
||
chunks = [content[i:i + MAX_LEN] for i in range(0, len(content), MAX_LEN)]
|
||
|
||
for chunk in chunks:
|
||
resp = requests.post(
|
||
webhook_url,
|
||
json={
|
||
"msgtype": "text",
|
||
"text": {"content": chunk},
|
||
},
|
||
timeout=10,
|
||
)
|
||
resp_data = resp.json()
|
||
if resp_data.get("errcode") != 0:
|
||
print(f"[WECOM ERROR] 群机器人发送失败: {resp_data}", flush=True)
|
||
else:
|
||
print(f"[WECOM DEBUG] 群机器人发送成功", flush=True)
|