220 lines
8.3 KiB
Python
220 lines
8.3 KiB
Python
|
from config import AgentId,Secret,corpid,token,encodingAESKey
|
|||
|
|
#import xml.etree.ElementTree as ET
|
|||
|
|
from flask import Flask,request
|
|||
|
|
from WXBizMsgCrypt import WXBizMsgCrypt
|
|||
|
|
import base64,hashlib
|
|||
|
|
from Crypto.Cipher import AES
|
|||
|
|
import time
|
|||
|
|
from xml.etree import ElementTree
|
|||
|
|
from keyword_config import KEYWORD_REPLIES
|
|||
|
|
from ai_service import OllamaService,MessageHandler,DifyService
|
|||
|
|
from typing import Dict
|
|||
|
|
from config import OPENAI_API_KEY, OPENAI_MODEL, OPENAI_BASE_URL,DIFY_API_KEY,DIFY_BASE_URL
|
|||
|
|
from config import FASTAPI_BASE_URL, FASTAPI_API_KEY
|
|||
|
|
import re
|
|||
|
|
|
|||
|
|
from ai_service import OpenAIService,FastGptService
|
|||
|
|
app =Flask(__name__)
|
|||
|
|
wxcpt = WXBizMsgCrypt(token,encodingAESKey,corpid)
|
|||
|
|
#ai_service = OllamaService()
|
|||
|
|
ai_service = OpenAIService(OPENAI_API_KEY, OPENAI_MODEL, OPENAI_BASE_URL)
|
|||
|
|
#ai_service = DifyService(DIFY_API_KEY,DIFY_BASE_URL)
|
|||
|
|
#ai_service = FastGptService(FASTAPI_API_KEY, FASTAPI_BASE_URL)
|
|||
|
|
|
|||
|
|
# 检查base64编码后数据位数是否正确
|
|||
|
|
def check_base64_len(base64_str):
|
|||
|
|
len_remainder = 4 - (len(base64_str) % 4)
|
|||
|
|
if len_remainder == 0:
|
|||
|
|
return base64_str
|
|||
|
|
else:
|
|||
|
|
for temp in range(0,len_remainder):
|
|||
|
|
base64_str = base64_str + "="
|
|||
|
|
return base64_str
|
|||
|
|
# 解密并提取消息正文
|
|||
|
|
def msg_base64_decrypt(ciphertext_base64,key_base64):
|
|||
|
|
# 处理密文、密钥和iv
|
|||
|
|
ciphertext_bytes = base64.b64decode(check_base64_len(ciphertext_base64))
|
|||
|
|
key_bytes = base64.b64decode(check_base64_len(key_base64))
|
|||
|
|
iv_bytes = key_bytes[:16]
|
|||
|
|
|
|||
|
|
# 解密
|
|||
|
|
decr = AES.new(key_bytes,AES.MODE_CBC,iv_bytes)
|
|||
|
|
plaintext_bytes = decr.decrypt(ciphertext_bytes)
|
|||
|
|
|
|||
|
|
# 截取数据，判断消息正文字节数
|
|||
|
|
msg_len_bytes = plaintext_bytes[16:20]
|
|||
|
|
msg_len = int.from_bytes(msg_len_bytes,byteorder='big', signed=False)
|
|||
|
|
|
|||
|
|
# 根据消息正文字节数截取消息正文，并转为字符串格式
|
|||
|
|
msg_bytes = plaintext_bytes[20:20+msg_len]
|
|||
|
|
msg = str(msg_bytes,encoding='utf-8')
|
|||
|
|
|
|||
|
|
return msg
|
|||
|
|
# 消息体签名校验
|
|||
|
|
def check_msg_signature(msg_signature,token,timestamp,nonce,echostr):
|
|||
|
|
# 使用sort()从小到大排序[].sort()是在原地址改值的，所以如果使用li_s = li.sort()，li_s是空的，li的值变为排序后的值]
|
|||
|
|
li = [token,timestamp,nonce,echostr]
|
|||
|
|
li.sort()
|
|||
|
|
# 将排序结果拼接
|
|||
|
|
li_str = li[0]+li[1]+li[2]+li[3]
|
|||
|
|
|
|||
|
|
# 计算SHA-1值
|
|||
|
|
sha1 = hashlib.sha1()
|
|||
|
|
# update()要指定加密字符串字符代码，不然要报错：
|
|||
|
|
# "Unicode-objects must be encoded before hashing"
|
|||
|
|
sha1.update(li_str.encode("utf8"))
|
|||
|
|
sha1_result = sha1.hexdigest()
|
|||
|
|
|
|||
|
|
# 比较并返回比较结果
|
|||
|
|
if sha1_result == msg_signature:
|
|||
|
|
return True
|
|||
|
|
else:
|
|||
|
|
return False
|
|||
|
|
@app.route('/hello', methods=['GET'])
|
|||
|
|
def hello():
|
|||
|
|
return "Hello, Flask!"
|
|||
|
|
|
|||
|
|
@app.route('/', methods=['GET', 'POST'])
|
|||
|
|
def reply():
|
|||
|
|
try:
|
|||
|
|
# 处理GET请求(验证URL)部分保持不变
|
|||
|
|
if request.method == 'GET':
|
|||
|
|
msg_signature = request.args.to_dict().get("msg_signature")
|
|||
|
|
timestamp = request.args.to_dict().get("timestamp")
|
|||
|
|
nonce = request.args.to_dict().get("nonce")
|
|||
|
|
echostr = request.args.to_dict().get("echostr")
|
|||
|
|
print(msg_signature,timestamp,nonce,echostr)
|
|||
|
|
# 获取消息体签名校验结果
|
|||
|
|
check_result = check_msg_signature(msg_signature,token,timestamp,nonce,echostr)
|
|||
|
|
if check_result:
|
|||
|
|
decrypt_result = msg_base64_decrypt(echostr,encodingAESKey)
|
|||
|
|
print("通过")
|
|||
|
|
return decrypt_result
|
|||
|
|
else:
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
# 处理POST请求(消息处理)
|
|||
|
|
elif request.method == 'POST':
|
|||
|
|
try:
|
|||
|
|
# 获取并解析原始数据部分保持不变
|
|||
|
|
raw_data = request.get_data()
|
|||
|
|
if isinstance(raw_data, bytes):
|
|||
|
|
raw_data = raw_data.decode('utf-8')
|
|||
|
|
|
|||
|
|
# 获取参数验证部分保持不变
|
|||
|
|
msg_signature = request.args.get('msg_signature', '')
|
|||
|
|
timestamp = str(request.args.get('timestamp', ''))
|
|||
|
|
nonce = request.args.get('nonce', '')
|
|||
|
|
|
|||
|
|
if not all([msg_signature, timestamp, nonce]):
|
|||
|
|
return "缺少必要参数", 400
|
|||
|
|
|
|||
|
|
# 解密消息部分保持不变
|
|||
|
|
ret, xml_content = wxcpt.DecryptMsg(
|
|||
|
|
raw_data,
|
|||
|
|
msg_signature,
|
|||
|
|
timestamp,
|
|||
|
|
nonce
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if ret != 0:
|
|||
|
|
print(f"[ERROR] 解密失败，错误码: {ret}")
|
|||
|
|
return "消息解密失败", 500
|
|||
|
|
|
|||
|
|
# 解析XML
|
|||
|
|
xml_tree = ElementTree.fromstring(xml_content)
|
|||
|
|
|
|||
|
|
def get_text(element):
|
|||
|
|
if element is None:
|
|||
|
|
return ''
|
|||
|
|
text = element.text
|
|||
|
|
if text is None:
|
|||
|
|
return ''
|
|||
|
|
return text.decode('utf-8') if isinstance(text, bytes) else str(text)
|
|||
|
|
|
|||
|
|
# 获取用户ID和消息内容
|
|||
|
|
from_user_name = get_text(xml_tree.find('FromUserName'))
|
|||
|
|
msg_content = get_text(xml_tree.find('Content'))
|
|||
|
|
|
|||
|
|
# 处理刷新对话的关键词
|
|||
|
|
refresh_keywords = ["new", "refresh", "00", "restart", "刷新", "新话题", "退下", "结束", "over"]
|
|||
|
|
if msg_content.strip().lower() in refresh_keywords:
|
|||
|
|
if hasattr(ai_service, 'conversation_history') and from_user_name in ai_service.conversation_history:
|
|||
|
|
del ai_service.conversation_history[from_user_name]
|
|||
|
|
re_text = "会话已重置"
|
|||
|
|
else:
|
|||
|
|
# 使用用户ID生成回复
|
|||
|
|
msg_handler = MessageHandler(KEYWORD_REPLIES, ai_service)
|
|||
|
|
re_text = msg_handler.get_reply(msg_content)
|
|||
|
|
if isinstance(ai_service, OpenAIService):
|
|||
|
|
re_text = ai_service.generate_response(msg_content, user_id=from_user_name)
|
|||
|
|
|
|||
|
|
# 处理回复文本
|
|||
|
|
re_text = process_text(re_text)
|
|||
|
|
|
|||
|
|
# 构造回复消息
|
|||
|
|
reply_dict = {
|
|||
|
|
'ToUserName': from_user_name,
|
|||
|
|
'FromUserName': get_text(xml_tree.find('ToUserName')),
|
|||
|
|
'CreateTime': str(int(time.time())),
|
|||
|
|
'MsgType': 'text',
|
|||
|
|
'Content': re_text
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 构造XML回复并加密
|
|||
|
|
reply_msg = ResponseMessage(reply_dict).xml
|
|||
|
|
|
|||
|
|
if isinstance(reply_msg, bytes):
|
|||
|
|
reply_msg = reply_msg.decode('utf-8')
|
|||
|
|
|
|||
|
|
ret, encrypt_xml = wxcpt.EncryptMsg(reply_msg, nonce, timestamp)
|
|||
|
|
|
|||
|
|
if ret != 0:
|
|||
|
|
print(f"[ERROR] 加密回复消息失败，错误码: {ret}")
|
|||
|
|
return "加密回复消息失败", 500
|
|||
|
|
|
|||
|
|
return encrypt_xml
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[ERROR] 处理消息时发生异常: {str(e)}")
|
|||
|
|
return "服务器内部错误", 500
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"处理异常: {str(e)}")
|
|||
|
|
return "服务器错误", 500
|
|||
|
|
|
|||
|
|
return "success"
|
|||
|
|
|
|||
|
|
# ResponseMessage类修改
|
|||
|
|
class ResponseMessage(object):
|
|||
|
|
def __init__(self, dict_data):
|
|||
|
|
self.dict_data = {k: str(v) for k, v in dict_data.items()}
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def xml(self):
|
|||
|
|
xml = "<xml>"
|
|||
|
|
for k, v in self.dict_data.items():
|
|||
|
|
xml += f"<{k}><![CDATA[{v}]]></{k}>"
|
|||
|
|
xml += "</xml>"
|
|||
|
|
return xml
|
|||
|
|
|
|||
|
|
|
|||
|
|
def process_text(text):
|
|||
|
|
# 1. 把所有的 "\n" 替换为 "__lineFeed__"
|
|||
|
|
text = text.replace('\n', '__lineFeed__')
|
|||
|
|
|
|||
|
|
# 2. 删除所有换行符（包括 \r 等）
|
|||
|
|
text = re.sub(r'[^\S]+', '', text)
|
|||
|
|
|
|||
|
|
# 3. 把所有的 "__lineFeed__" 替换回 "\n"
|
|||
|
|
text = text.replace('__lineFeed__', '\n')
|
|||
|
|
# 4. 删除多余的换行符
|
|||
|
|
text = re.sub(r'\n+', '\n', text)
|
|||
|
|
return text.strip()
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
app.run(host='0.0.0.0', port=5000, debug=True)
|