187 lines
7.4 KiB
Python
187 lines
7.4 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
企业微信会话存档SDK命令行工具
|
||
Python重构版本，替代原始的tool_testSdk.cpp
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import os
|
||
from pathlib import Path
|
||
from dotenv import load_dotenv
|
||
|
||
# 加载环境变量
|
||
load_dotenv()
|
||
|
||
# 添加父目录到路径，以便导入wework_ctypes
|
||
sys.path.insert(0, str(Path(__file__).parent))
|
||
|
||
from wework_ctypes import WeWorkFinanceSDK, WeWorkSDKError
|
||
|
||
def parse_args():
|
||
"""解析命令行参数"""
|
||
# 从环境变量获取默认值
|
||
default_corpid = os.getenv("CORPID", "wwdf65802ca25ec195")
|
||
default_secret = os.getenv("WEWORK_SECRET", "-Ta6WMWxBhfGolWnnlO15nQckj3DRKAowUOdX2fwvzE")
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="企业微信会话存档SDK工具",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
使用示例:
|
||
# 拉取会话存档
|
||
%(prog)s chatmsg 0 100 "" "" 5
|
||
|
||
# 拉取媒体文件
|
||
%(prog)s mediadata "fileid" "" "" 5 "/tmp/media.file"
|
||
|
||
# 解密数据
|
||
%(prog)s decryptdata "encrypt_key" "encrypt_chat_msg"
|
||
|
||
# 使用代理拉取会话存档
|
||
%(prog)s chatmsg 0 100 "http://proxy.example.com:8080" "user:pass" 10
|
||
"""
|
||
)
|
||
|
||
subparsers = parser.add_subparsers(dest='command', help='命令')
|
||
|
||
# chatmsg命令
|
||
chatmsg_parser = subparsers.add_parser('chatmsg', help='拉取会话存档')
|
||
chatmsg_parser.add_argument('seq', type=int, help='从指定的seq开始拉取消息，首次使用请使用0')
|
||
chatmsg_parser.add_argument('limit', type=int, help='一次拉取的消息条数，最大值1000条')
|
||
chatmsg_parser.add_argument('proxy', type=str, help='代理链接，如：socks5://10.0.0.1:8081 或 http://10.0.0.1:8081')
|
||
chatmsg_parser.add_argument('passwd', type=str, help='代理账号密码，如：user_name:passwd_123')
|
||
chatmsg_parser.add_argument('timeout', type=int, help='超时时间，单位秒')
|
||
|
||
# mediadata命令
|
||
mediadata_parser = subparsers.add_parser('mediadata', help='拉取媒体文件')
|
||
mediadata_parser.add_argument('fileid', type=str, help='媒体文件id，从解密后的会话存档中得到')
|
||
mediadata_parser.add_argument('proxy', type=str, help='代理链接')
|
||
mediadata_parser.add_argument('passwd', type=str, help='代理账号密码')
|
||
mediadata_parser.add_argument('timeout', type=int, help='超时时间，单位秒')
|
||
mediadata_parser.add_argument('savefile', type=str, help='媒体文件保存路径')
|
||
|
||
# decryptdata命令
|
||
decryptdata_parser = subparsers.add_parser('decryptdata', help='解密会话存档内容')
|
||
decryptdata_parser.add_argument('encrypt_key', type=str, help='拉取会话存档返回的encrypt_random_key，使用配置在企业微信管理台的rsa公钥对应的私钥解密后得到encrypt_key')
|
||
decryptdata_parser.add_argument('encrypt_chat_msg', type=str, help='拉取会话存档返回的encrypt_chat_msg')
|
||
|
||
# 添加全局选项
|
||
parser.add_argument('--corpid', type=str,
|
||
default=default_corpid,
|
||
help=f'企业ID，默认从环境变量CORPID读取 (当前: {default_corpid})')
|
||
parser.add_argument('--secret', type=str,
|
||
default=default_secret,
|
||
help=f'聊天内容存档的Secret，默认从环境变量WEWORK_SECRET读取')
|
||
parser.add_argument('--lib-path', type=str,
|
||
help='libWeWorkFinanceSdk_C.so路径，默认使用lib/sdk/libWeWorkFinanceSdk_C.so')
|
||
|
||
return parser.parse_args()
|
||
|
||
def main():
|
||
"""主函数"""
|
||
args = parse_args()
|
||
|
||
if not args.command:
|
||
print("错误: 需要指定命令 (chatmsg, mediadata, decryptdata)")
|
||
print("使用 --help 查看帮助")
|
||
return 1
|
||
|
||
try:
|
||
# 初始化SDK
|
||
with WeWorkFinanceSDK(args.corpid, args.secret, args.lib_path) as sdk:
|
||
|
||
if args.command == 'chatmsg':
|
||
# 拉取会话存档
|
||
print(f"拉取会话存档: seq={args.seq}, limit={args.limit}, timeout={args.timeout}s")
|
||
if args.proxy:
|
||
print(f"使用代理: {args.proxy}")
|
||
|
||
result = sdk.get_chat_data(
|
||
seq=args.seq,
|
||
limit=args.limit,
|
||
proxy=args.proxy,
|
||
passwd=args.passwd,
|
||
timeout=args.timeout
|
||
)
|
||
|
||
# 尝试解析JSON并美化输出
|
||
try:
|
||
data = json.loads(result)
|
||
print("获取聊天数据成功:")
|
||
print(json.dumps(data, indent=2, ensure_ascii=False))
|
||
|
||
# 如果有聊天数据，显示统计信息
|
||
if 'chatdata' in data and data['chatdata']:
|
||
print(f"\n共获取 {len(data['chatdata'])} 条聊天记录")
|
||
if data['chatdata']:
|
||
max_seq = max(msg.get('seq', 0) for msg in data['chatdata'])
|
||
print(f"最大seq值: {max_seq} (下次拉取建议使用此值)")
|
||
|
||
except json.JSONDecodeError:
|
||
print("获取聊天数据成功 (原始数据):")
|
||
print(result)
|
||
|
||
return 0
|
||
|
||
elif args.command == 'mediadata':
|
||
# 拉取媒体文件
|
||
print(f"拉取媒体文件: fileid={args.fileid}, timeout={args.timeout}s")
|
||
print(f"保存到: {args.savefile}")
|
||
if args.proxy:
|
||
print(f"使用代理: {args.proxy}")
|
||
|
||
success = sdk.get_media_data(
|
||
sdk_fileid=args.fileid,
|
||
save_path=args.savefile,
|
||
proxy=args.proxy,
|
||
passwd=args.passwd,
|
||
timeout=args.timeout
|
||
)
|
||
|
||
if success:
|
||
print(f"媒体文件下载完成: {args.savefile}")
|
||
|
||
# 检查文件大小
|
||
if os.path.exists(args.savefile):
|
||
size = os.path.getsize(args.savefile)
|
||
print(f"文件大小: {size} 字节 ({size/1024:.2f} KB)")
|
||
else:
|
||
print("媒体文件下载失败")
|
||
return 1
|
||
|
||
return 0
|
||
|
||
elif args.command == 'decryptdata':
|
||
# 解密会话存档内容
|
||
print("解密会话存档内容...")
|
||
|
||
result = sdk.decrypt_data(args.encrypt_key, args.encrypt_chat_msg)
|
||
|
||
# 尝试解析JSON并美化输出
|
||
try:
|
||
data = json.loads(result)
|
||
print("解密成功:")
|
||
print(json.dumps(data, indent=2, ensure_ascii=False))
|
||
except json.JSONDecodeError:
|
||
print("解密成功 (原始数据):")
|
||
print(result)
|
||
|
||
return 0
|
||
|
||
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
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|