Files
chat_app/backend/scripts/load_chat_data.py
2025-12-11 03:53:47 +08:00

705 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
企业微信会话存档SDK测试脚本
简单测试脚本,调用SDK获取会话数据并自动解密和存储到数据库
使用示例:
python test_chat_data.py # 使用默认seq=0
python test_chat_data.py --seq 1000 # 从seq=1000开始拉取
python test_chat_data.py -s 5000 # 从seq=5000开始拉取
python test_chat_data.py --help # 显示帮助信息
"""
import argparse
import asyncio
import json
import os
import sys
import base64
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()
# 添加当前目录到路径,以便导入wework_ctypes
sys.path.insert(0, str(Path(__file__).parent))
# 添加backend目录到路径,以便导入模型
sys.path.insert(0, str(Path(__file__).parent.parent))
from wework_ctypes import WeWorkFinanceSDK, WeWorkSDKError, ERROR_CODES
def load_rsa_private_key(key_path: str) -> str:
"""
加载RSA私钥
Args:
key_path: 私钥文件路径
Returns:
私钥内容字符串
"""
try:
with open(key_path, "r") as f:
key_content = f.read()
return key_content
except Exception as e:
raise Exception(f"加载RSA私钥失败: {e}")
def rsa_decrypt_base64(encrypted_data_base64: str, private_key: str) -> str:
"""
使用RSA私钥解密Base64编码的数据
Args:
encrypted_data_base64: Base64编码的加密数据
private_key: RSA私钥字符串
Returns:
解密后的字符串
"""
try:
# 使用cryptography库进行RSA解密
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.backends import default_backend
import base64
# 加载私钥
private_key_obj = serialization.load_pem_private_key(
private_key.encode("utf-8"), password=None, backend=default_backend()
)
# Base64解码
encrypted_data = base64.b64decode(encrypted_data_base64)
# RSA解密 (使用PKCS1v15填充)
decrypted_data = private_key_obj.decrypt(encrypted_data, padding.PKCS1v15())
return decrypted_data.decode("utf-8")
except ImportError:
raise Exception("需要安装cryptography库: pip install cryptography")
except Exception as e:
raise Exception(f"RSA解密失败: {e}")
def decrypt_chat_message(
sdk,
encrypt_random_key: str,
encrypt_chat_msg: str,
rsa_private_key_path: str,
publickey_ver: int = None,
) -> dict:
"""
解密聊天消息
Args:
sdk: WeWorkFinanceSDK实例
encrypt_random_key: 加密的随机密钥(Base64编码)
encrypt_chat_msg: 加密的聊天消息
rsa_private_key_path: RSA私钥文件路径或目录
publickey_ver: publickey_ver版本号,用于动态匹配私钥
Returns:
解密后的消息字典
"""
# 动态确定私钥路径
actual_key_path = rsa_private_key_path
# 检查是否使用默认私钥路径(用户未指定自定义路径)
# 通过检查路径是否指向默认文件来判断
default_key_path = Path("cert/1_rsa.key")
specified_key_path = Path(rsa_private_key_path)
# 如果使用默认私钥路径且提供了publickey_ver,则动态匹配
if publickey_ver is not None and specified_key_path == default_key_path:
# 构建动态私钥路径:cert/{publickey_ver}_rsa.key
dynamic_key_path = f"cert/{publickey_ver}_rsa.key"
dynamic_path = Path(dynamic_key_path)
if dynamic_path.exists():
actual_key_path = dynamic_key_path
else:
# 如果动态私钥不存在,使用默认私钥
pass
# 1. 加载RSA私钥
private_key = load_rsa_private_key(actual_key_path)
# 2. 使用RSA私钥解密encrypt_random_key得到encrypt_key
try:
encrypt_key = rsa_decrypt_base64(encrypt_random_key, private_key)
except Exception as e:
# 尝试直接使用encrypt_random_key作为encrypt_key(某些情况下可能已经解密)
encrypt_key = encrypt_random_key
# 3. 使用SDK解密聊天消息
decrypted_msg = sdk.decrypt_data(encrypt_key, encrypt_chat_msg)
# 4. 解析JSON
try:
result = json.loads(decrypted_msg)
return result
except json.JSONDecodeError:
# 如果不是JSON,返回原始文本
return {"raw_message": decrypted_msg}
def decrypt_chat_data(chatdata, sdk, rsa_private_key_path="cert/1_rsa.key"):
"""
解密聊天数据
Args:
chatdata: 原始聊天数据列表
sdk: WeWorkFinanceSDK实例
rsa_private_key_path: RSA私钥文件路径
Returns:
(decrypted_chatdata, stats)
decrypted_chatdata: 解密后的聊天数据列表
stats: 解密统计信息字典
"""
decrypted_chatdata = []
stats = {
"total_encrypted_messages": 0,
"successful_decrypts": 0,
"failed_decrypts": 0,
}
# 检查私钥文件是否存在
rsa_path = Path(rsa_private_key_path)
if not rsa_path.exists():
print(f"警告: RSA私钥文件不存在: {rsa_private_key_path}")
print("将跳过解密,保存原始加密数据")
return chatdata, stats
print(f"使用RSA私钥: {rsa_path.absolute()}")
print("开始解密聊天数据...")
for i, msg in enumerate(chatdata):
encrypt_random_key = msg.get("encrypt_random_key", "")
encrypt_chat_msg = msg.get("encrypt_chat_msg", "")
# 检查是否有加密数据
if not encrypt_random_key or not encrypt_chat_msg:
# 没有加密数据,直接添加到结果中
decrypted_chatdata.append(msg)
continue
stats["total_encrypted_messages"] += 1
try:
# 获取publickey_ver
publickey_ver = msg.get("publickey_ver")
# 解密消息
decrypted_msg = decrypt_chat_message(
sdk,
encrypt_random_key,
encrypt_chat_msg,
rsa_private_key_path,
publickey_ver,
)
# 创建新的消息对象,保留原始字段但移除加密字段
new_msg = {}
for key, value in msg.items():
if key not in ["encrypt_random_key", "encrypt_chat_msg"]:
new_msg[key] = value
# 将解密后的字段合并到消息中
if isinstance(decrypted_msg, dict):
for key, value in decrypted_msg.items():
new_msg[key] = value
decrypted_chatdata.append(new_msg)
stats["successful_decrypts"] += 1
# 显示进度
if stats["successful_decrypts"] % 100 == 0:
print(f" 已解密 {stats['successful_decrypts']} 条消息...")
except Exception as e:
stats["failed_decrypts"] += 1
# 解密失败,跳过此消息(不添加到结果中)
continue
return decrypted_chatdata, stats
async def save_messages_to_database(messages: List[Dict[str, Any]]) -> Dict[str, int]:
"""
将消息保存到数据库
Args:
messages: 消息字典列表
Returns:
存储统计信息字典
"""
stats = {
"total_messages": len(messages),
"saved": 0,
"skipped_duplicates": 0,
"failed": 0,
}
if not messages:
print("没有消息需要存储到数据库")
return stats
print(f"\n开始将 {len(messages)} 条消息存储到数据库...")
try:
# 导入数据库相关模块
from pymongo import AsyncMongoClient
from beanie import init_beanie
from app.models import models
from app.models.message import Message
# 从环境变量获取数据库配置
mongodb_url = os.getenv(
"MONGODB_URL",
"mongodb://localhost:27017",
)
database_name = os.getenv("DATABASE_NAME", "douba_chat_dev")
# 初始化数据库连接
client = AsyncMongoClient(mongodb_url)
await init_beanie(database=client[database_name], document_models=models)
print(f"数据库连接成功: {database_name}")
# 批量处理消息
batch_size = 100
for i in range(0, len(messages), batch_size):
batch = messages[i : i + batch_size]
batch_msgids = [msg.get("msgid") for msg in batch if msg.get("msgid")]
# 检查哪些消息已经存在
existing_messages = await Message.find(
{"msgid": {"$in": batch_msgids}}
).to_list()
existing_msgids = {msg.msgid for msg in existing_messages}
# 过滤掉已存在的消息
new_messages = []
for msg in batch:
msgid = msg.get("msgid")
if not msgid:
stats["failed"] += 1
continue
if msgid in existing_msgids:
stats["skipped_duplicates"] += 1
continue
try:
# 创建Message实例
# 注意:使用by_alias=True来处理"from"字段的别名
message_data = {
"msgid": msg.get("msgid"),
"action": msg.get("action"),
"from": msg.get("from"), # 使用别名
"tolist": msg.get("tolist", []),
"roomid": msg.get("roomid"),
"text": msg.get("text", {}),
"msgtime": msg.get("msgtime"),
"msgtype": msg.get("msgtype"),
"remark": msg.get("remark"),
"seq": msg.get("seq", 1000),
"isEnabled": True,
"created_at": datetime.now(),
"updated_at": datetime.now(),
}
# 移除None值
message_data = {
k: v for k, v in message_data.items() if v is not None
}
message = Message(**message_data)
new_messages.append(message)
except Exception as e:
print(f"创建消息实例失败 (msgid: {msgid}): {e}")
stats["failed"] += 1
# 批量插入新消息
if new_messages:
try:
await Message.insert_many(new_messages)
stats["saved"] += len(new_messages)
print(
f" 已存储 {len(new_messages)} 条消息 (进度: {i + len(batch)}/{len(messages)})"
)
except Exception as e:
print(f"批量插入失败: {e}")
stats["failed"] += len(new_messages)
# 关闭数据库连接
await client.close()
print(f"数据库存储完成:")
print(f" 总计消息: {stats['total_messages']}")
print(f" 成功存储: {stats['saved']}")
print(f" 跳过重复: {stats['skipped_duplicates']}")
print(f" 存储失败: {stats['failed']}")
except Exception as e:
print(f"数据库存储过程中发生错误: {e}")
import traceback
traceback.print_exc()
stats["failed"] = len(messages)
return stats
def main():
"""主函数:测试获取会话数据"""
# 解析命令行参数
parser = argparse.ArgumentParser(
description="企业微信会话存档SDK测试脚本",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
使用示例:
%(prog)s # 使用默认seq=0
%(prog)s --seq 1000 # 从seq=1000开始拉取
%(prog)s -s 5000 # 从seq=5000开始拉取
%(prog)s --limit 500 # 一次拉取500条消息
%(prog)s --timeout 30 # 设置超时时间为30秒
""",
)
parser.add_argument(
"-s",
"--seq",
type=int,
default=0,
help="起始seq值,首次拉取使用0,后续拉取使用上次返回的最大seq值 (默认: 0)",
)
parser.add_argument(
"-l",
"--limit",
type=int,
default=1000,
help="一次拉取的消息数量限制 (默认: 1000)",
)
parser.add_argument(
"-t", "--timeout", type=int, default=10, help="请求超时时间(秒) (默认: 10)"
)
parser.add_argument(
"--proxy",
type=str,
default="",
help="代理服务器地址,格式: http://proxy.example.com:8080 (默认: 空)",
)
parser.add_argument(
"--passwd", type=str, default="", help="代理服务器密码 (默认: 空)"
)
args = parser.parse_args()
print("=" * 60)
print("企业微信会话存档SDK测试脚本")
print("=" * 60)
print(f"命令行参数: seq={args.seq}, limit={args.limit}, timeout={args.timeout}s")
if args.proxy:
print(f"代理设置: {args.proxy}")
# 从环境变量获取配置
corpid = os.getenv("CORPID", "wwdf65802ca25ec195")
secret = os.getenv("WEWORK_SECRET", "-Ta6WMWxBhfGolWnnlO15nQckj3DRKAowUOdX2fwvzE")
print(f"企业ID: {corpid}")
print(f"Secret: {secret[:10]}...") # 只显示前10个字符,避免泄露完整密钥
# 检查动态库是否存在
lib_path = Path(__file__).parent.parent / "lib" / "sdk" / "libWeWorkFinanceSdk_C.so"
if not lib_path.exists():
print(f"错误: 动态库不存在: {lib_path}")
print("请确保已正确放置libWeWorkFinanceSdk_C.so文件")
return 1
print(f"动态库路径: {lib_path}")
# 确保data目录存在
data_dir = Path("data")
data_dir.mkdir(exist_ok=True)
print(f"数据目录: {data_dir.absolute()}")
print("-" * 60)
try:
# 初始化SDK(使用上下文管理器确保资源正确释放)
print("正在初始化SDK...")
with WeWorkFinanceSDK(corpid, secret, str(lib_path)) as sdk:
print("✓ SDK初始化成功")
# 设置拉取参数(从命令行参数获取)
seq = args.seq
limit = args.limit
timeout = args.timeout
print(f"\n正在拉取会话数据...")
print(f"参数: seq={seq}, limit={limit}, timeout={timeout}s")
if args.proxy:
print(f"代理: {args.proxy}")
# 调用SDK获取聊天数据
result_str = sdk.get_chat_data(
seq=seq,
limit=limit,
proxy=args.proxy, # 使用命令行指定的代理
passwd=args.passwd, # 使用命令行指定的代理密码
timeout=timeout,
)
print("✓ 数据拉取成功")
print("-" * 60)
# 尝试解析JSON
try:
result = json.loads(result_str)
# 检查错误码
errcode = result.get("errcode", -1)
errmsg = result.get("errmsg", "")
if errcode == 0:
print(f"成功: {errmsg}")
# 获取聊天数据
chatdata = result.get("chatdata", [])
original_chat_count = len(chatdata)
print(f"共获取 {original_chat_count} 条聊天记录")
if original_chat_count > 0:
# 显示第一条消息的摘要信息
print("\n第一条消息摘要:")
first_msg = chatdata[0]
print(f" seq: {first_msg.get('seq', 'N/A')}")
print(f" msgid: {first_msg.get('msgid', 'N/A')}")
print(
f" publickey_ver: {first_msg.get('publickey_ver', 'N/A')}"
)
# 检查是否有加密数据
encrypt_random_key = first_msg.get("encrypt_random_key", "")
encrypt_chat_msg = first_msg.get("encrypt_chat_msg", "")
if encrypt_random_key and encrypt_chat_msg:
print(f" 加密密钥长度: {len(encrypt_random_key)}")
print(f" 加密消息长度: {len(encrypt_chat_msg)}")
print(" 注意: 消息已加密,将自动解密")
else:
print(" 消息未加密")
# 解密聊天数据
decrypted_chatdata, decrypt_stats = decrypt_chat_data(chatdata, sdk)
# 更新结果中的chatdata
result["chatdata"] = decrypted_chatdata
decrypted_chat_count = len(decrypted_chatdata)
# 显示解密统计
print("\n" + "-" * 60)
print("解密完成:")
print(f" 原始消息总数: {original_chat_count}")
print(
f" 加密消息数量: {decrypt_stats['total_encrypted_messages']}"
)
print(f" 成功解密: {decrypt_stats['successful_decrypts']}")
print(f" 解密失败: {decrypt_stats['failed_decrypts']}")
print(f" 最终输出消息数: {decrypted_chat_count}")
# 显示最大seq值,用于下次拉取
if decrypted_chat_count > 0:
max_seq = max(msg.get("seq", 0) for msg in decrypted_chatdata)
print(f"\n最大seq值: {max_seq} (下次拉取建议使用此值)")
else:
max_seq = seq
# 生成文件名:chat_{start_seq}_{end_seq}_{date}.json
start_seq = seq
end_seq = max_seq if decrypted_chat_count > 0 else seq
date_str = datetime.now().strftime("%y%m%d") # YYMMDD格式
filename = f"chat_{start_seq}_{end_seq}_{date_str}.json"
output_path = data_dir / filename
# 添加元数据到结果中
if "_metadata" not in result:
result["_metadata"] = {}
# 更新chat_count为解密后的消息数量
result["_metadata"].update(
{
"filename": filename,
"save_path": str(output_path),
"save_time": datetime.now().isoformat(),
"seq_range": f"{start_seq}-{end_seq}",
"chat_count": decrypted_chat_count, # 使用解密后的数量
"parameters": {
"request_seq": seq,
"limit": limit,
"timeout": timeout,
"proxy": args.proxy if args.proxy else "未使用",
"command_line_args": {
"seq": args.seq,
"limit": args.limit,
"timeout": args.timeout,
"proxy": args.proxy,
"passwd": "已设置" if args.passwd else "未设置",
},
},
"decrypt_info": {
"total_encrypted_messages": decrypt_stats[
"total_encrypted_messages"
],
"successful_decrypts": decrypt_stats[
"successful_decrypts"
],
"failed_decrypts": decrypt_stats["failed_decrypts"],
"decrypt_time": datetime.now().isoformat(),
},
}
)
# 保存到data目录
with open(output_path, "w", encoding="utf-8") as f:
json.dump(result, f, indent=2, ensure_ascii=False)
print(f"\n解密结果已保存到: {output_path}")
print(f"文件名格式: chat_{start_seq}_{end_seq}_{date_str}.json")
print(
f"说明: seq范围 {start_seq}-{end_seq}, 日期 {datetime.now().strftime('%Y-%m-%d')}"
)
print(f"包含 {decrypted_chat_count} 条解密后的消息")
# 显示一些解密后的消息示例
if decrypted_chat_count > 0:
print("\n解密消息示例:")
for i, msg in enumerate(decrypted_chatdata[:3]): # 显示前3条
seq = msg.get("seq", "N/A")
msgid = msg.get("msgid", "N/A")[:20]
msgtype = msg.get("msgtype", "unknown")
from_user = msg.get("from", "unknown")
print(
f" [{i+1}] seq={seq}, msgid={msgid}..., 类型={msgtype}, 来自={from_user}"
)
# 存储消息到数据库
if decrypted_chat_count > 0:
print("\n" + "=" * 60)
print("开始存储消息到数据库...")
try:
# 运行异步函数
db_stats = asyncio.run(
save_messages_to_database(decrypted_chatdata)
)
# 更新元数据中的数据库存储信息
if "db_store_info" not in result["_metadata"]:
result["_metadata"]["db_store_info"] = {}
result["_metadata"]["db_store_info"].update(
{
"db_stats": db_stats,
"db_store_time": datetime.now().isoformat(),
}
)
# 重新保存JSON文件以包含数据库存储信息
with open(output_path, "w", encoding="utf-8") as f:
json.dump(result, f, indent=2, ensure_ascii=False)
print(f"数据库存储信息已更新到: {output_path}")
except Exception as e:
print(f"数据库存储过程中发生错误: {e}")
import traceback
traceback.print_exc()
else:
print(f"错误: errcode={errcode}, errmsg={errmsg}")
# 显示详细的错误说明
if errcode in ERROR_CODES:
print(f"详细说明: {ERROR_CODES[errcode]}")
# 根据错误码提供建议
if errcode == 10000:
print("建议: 请检查请求参数是否正确")
elif errcode == 10001:
print("建议: 请检查网络连接和代理设置")
elif errcode == 10002:
print("建议: 数据格式可能不正确,请检查响应数据")
elif errcode == 10004:
print("建议: 请检查密钥配置是否正确")
elif errcode == 10006:
print("建议: 解密失败,请检查私钥和加密数据")
elif errcode == 10007:
print("建议: 找不到消息加密版本的私钥,需要重新传入私钥对")
elif errcode == 10009:
print("建议: IP地址不在白名单中,请检查IP配置")
elif errcode == 10010:
print("建议: 数据已过期,请重新拉取数据")
elif errcode == 10011:
print("建议: 证书错误,请检查证书配置")
print("\n完整响应:")
print(json.dumps(result, indent=2, ensure_ascii=False))
except json.JSONDecodeError as e:
print(f"JSON解析错误: {e}")
print("原始响应数据:")
print(result_str[:500] + "..." if len(result_str) > 500 else result_str)
# 保存原始数据到文件(使用日期命名)
date_str = datetime.now().strftime("%y%m%d")
raw_filename = f"chat_raw_{date_str}.txt"
raw_output_path = data_dir / raw_filename
with open(raw_output_path, "w", encoding="utf-8") as f:
f.write(result_str)
print(f"原始数据已保存到: {raw_output_path}")
# 同时保存一份到当前目录
with open("chat_data_raw.txt", "w", encoding="utf-8") as f:
f.write(result_str)
print(f"备份原始数据: chat_data_raw.txt")
print("\n" + "=" * 60)
print("测试完成")
except WeWorkSDKError as e:
print(f"SDK错误: {e}")
return 1
except FileNotFoundError as e:
print(f"文件错误: {e}")
return 1
except Exception as e:
print(f"未知错误: {e}")
import traceback
traceback.print_exc()
return 1
return 0
if __name__ == "__main__":
sys.exit(main())