311 lines
10 KiB
Python
311 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
企业微信会话存档服务
|
||
集成到FastAPI项目中
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Optional, Dict, Any, List
|
||
from datetime import datetime
|
||
|
||
# 添加scripts目录到路径
|
||
import sys
|
||
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
|
||
|
||
from wework_ctypes import WeWorkFinanceSDK, WeWorkSDKError
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class WeworkService:
|
||
"""企业微信会话存档服务"""
|
||
|
||
def __init__(self, corpid: Optional[str] = None, secret: Optional[str] = None):
|
||
"""
|
||
初始化企业微信服务
|
||
|
||
Args:
|
||
corpid: 企业ID，如果为None则从环境变量CORPID读取
|
||
secret: 聊天内容存档的Secret，如果为None则从环境变量WEWORK_SECRET读取
|
||
"""
|
||
self.corpid = corpid or os.getenv("CORPID", "wwdf65802ca25ec195")
|
||
self.secret = secret or os.getenv("WEWORK_SECRET", "-Ta6WMWxBhfGolWnnlO15nQckj3DRKAowUOdX2fwvzE")
|
||
|
||
# 媒体文件保存目录
|
||
self.media_dir = Path(os.getenv("WEWORK_MEDIA_DIR", "/tmp/wework_media"))
|
||
self.media_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# SDK实例（延迟初始化）
|
||
self._sdk: Optional[WeWorkFinanceSDK] = None
|
||
|
||
logger.info(f"企业微信服务初始化完成，企业ID: {self.corpid}")
|
||
logger.info(f"媒体文件保存目录: {self.media_dir}")
|
||
|
||
@property
|
||
def sdk(self) -> WeWorkFinanceSDK:
|
||
"""获取SDK实例（延迟初始化）"""
|
||
if self._sdk is None:
|
||
try:
|
||
self._sdk = WeWorkFinanceSDK(self.corpid, self.secret)
|
||
logger.info("企业微信SDK初始化成功")
|
||
except Exception as e:
|
||
logger.error(f"企业微信SDK初始化失败: {e}")
|
||
raise
|
||
return self._sdk
|
||
|
||
def get_chat_data(self, seq: int = 0, limit: int = 100,
|
||
proxy: str = "", passwd: str = "",
|
||
timeout: int = 5) -> Dict[str, Any]:
|
||
"""
|
||
拉取会话存档
|
||
|
||
Args:
|
||
seq: 从指定的seq开始拉取消息，首次使用请使用0
|
||
limit: 一次拉取的消息条数，最大值1000条
|
||
proxy: 代理链接
|
||
passwd: 代理账号密码
|
||
timeout: 超时时间，单位秒
|
||
|
||
Returns:
|
||
包含聊天数据的字典
|
||
|
||
Raises:
|
||
WeWorkSDKError: SDK操作失败
|
||
"""
|
||
logger.info(f"拉取会话存档: seq={seq}, limit={limit}, timeout={timeout}s")
|
||
|
||
try:
|
||
result_str = self.sdk.get_chat_data(seq, limit, proxy, passwd, timeout)
|
||
|
||
# 解析JSON结果
|
||
result = json.loads(result_str)
|
||
|
||
# 添加元数据
|
||
result["_metadata"] = {
|
||
"timestamp": datetime.now().isoformat(),
|
||
"params": {
|
||
"seq": seq,
|
||
"limit": limit,
|
||
"timeout": timeout,
|
||
"has_proxy": bool(proxy)
|
||
}
|
||
}
|
||
|
||
# 如果有聊天数据，记录统计信息
|
||
if "chatdata" in result and result["chatdata"]:
|
||
chat_count = len(result["chatdata"])
|
||
max_seq = max(msg.get("seq", 0) for msg in result["chatdata"])
|
||
logger.info(f"成功拉取 {chat_count} 条聊天记录，最大seq: {max_seq}")
|
||
|
||
# 提取消息类型统计
|
||
msg_types = {}
|
||
for msg in result["chatdata"]:
|
||
# 这里需要根据实际消息格式解析消息类型
|
||
pass
|
||
|
||
return result
|
||
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"解析聊天数据JSON失败: {e}")
|
||
raise WeWorkSDKError(-1, f"解析聊天数据失败: {e}")
|
||
|
||
def get_media_data(self, sdk_fileid: str, filename: Optional[str] = None,
|
||
proxy: str = "", passwd: str = "",
|
||
timeout: int = 5) -> Dict[str, Any]:
|
||
"""
|
||
拉取媒体文件
|
||
|
||
Args:
|
||
sdk_fileid: 媒体文件id
|
||
filename: 保存的文件名，如果为None则自动生成
|
||
proxy: 代理链接
|
||
passwd: 代理账号密码
|
||
timeout: 超时时间，单位秒
|
||
|
||
Returns:
|
||
包含文件信息的字典
|
||
|
||
Raises:
|
||
WeWorkSDKError: SDK操作失败
|
||
"""
|
||
logger.info(f"拉取媒体文件: fileid={sdk_fileid}, timeout={timeout}s")
|
||
|
||
# 生成文件名
|
||
if filename is None:
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
filename = f"media_{timestamp}_{sdk_fileid[:8]}.bin"
|
||
|
||
save_path = self.media_dir / filename
|
||
|
||
try:
|
||
success = self.sdk.get_media_data(
|
||
sdk_fileid=sdk_fileid,
|
||
save_path=str(save_path),
|
||
proxy=proxy,
|
||
passwd=passwd,
|
||
timeout=timeout
|
||
)
|
||
|
||
if success and save_path.exists():
|
||
file_size = save_path.stat().st_size
|
||
logger.info(f"媒体文件下载完成: {filename}, 大小: {file_size} 字节")
|
||
|
||
return {
|
||
"success": True,
|
||
"filename": filename,
|
||
"filepath": str(save_path),
|
||
"filesize": file_size,
|
||
"download_time": datetime.now().isoformat(),
|
||
"fileid": sdk_fileid
|
||
}
|
||
else:
|
||
logger.error(f"媒体文件下载失败: {filename}")
|
||
return {
|
||
"success": False,
|
||
"error": "下载失败",
|
||
"fileid": sdk_fileid
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"拉取媒体文件时出错: {e}")
|
||
raise
|
||
|
||
def decrypt_data(self, encrypt_key: str, encrypt_msg: str) -> Dict[str, Any]:
|
||
"""
|
||
解密会话内容
|
||
|
||
Args:
|
||
encrypt_key: 解密密钥
|
||
encrypt_msg: 加密的消息
|
||
|
||
Returns:
|
||
包含解密数据的字典
|
||
|
||
Raises:
|
||
WeWorkSDKError: SDK操作失败
|
||
"""
|
||
logger.info("解密会话存档内容")
|
||
|
||
try:
|
||
result_str = self.sdk.decrypt_data(encrypt_key, encrypt_msg)
|
||
|
||
# 解析JSON结果
|
||
result = json.loads(result_str)
|
||
|
||
# 添加元数据
|
||
result["_metadata"] = {
|
||
"timestamp": datetime.now().isoformat(),
|
||
"decrypted": True
|
||
}
|
||
|
||
logger.info("解密成功")
|
||
return result
|
||
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"解析解密数据JSON失败: {e}")
|
||
# 如果不是JSON，返回原始文本
|
||
return {
|
||
"success": True,
|
||
"decrypted_data": result_str,
|
||
"is_json": False,
|
||
"_metadata": {
|
||
"timestamp": datetime.now().isoformat(),
|
||
"decrypted": True
|
||
}
|
||
}
|
||
|
||
def get_media_files(self) -> List[Dict[str, Any]]:
|
||
"""
|
||
获取已下载的媒体文件列表
|
||
|
||
Returns:
|
||
媒体文件信息列表
|
||
"""
|
||
files = []
|
||
for file_path in self.media_dir.glob("*"):
|
||
if file_path.is_file():
|
||
stat = file_path.stat()
|
||
files.append({
|
||
"filename": file_path.name,
|
||
"filepath": str(file_path),
|
||
"filesize": stat.st_size,
|
||
"created": datetime.fromtimestamp(stat.st_ctime).isoformat(),
|
||
"modified": datetime.fromtimestamp(stat.st_mtime).isoformat()
|
||
})
|
||
|
||
return sorted(files, key=lambda x: x["modified"], reverse=True)
|
||
|
||
def cleanup_old_files(self, days: int = 7) -> int:
|
||
"""
|
||
清理旧的媒体文件
|
||
|
||
Args:
|
||
days: 保留天数
|
||
|
||
Returns:
|
||
删除的文件数量
|
||
"""
|
||
cutoff_time = datetime.now().timestamp() - (days * 24 * 3600)
|
||
deleted_count = 0
|
||
|
||
for file_path in self.media_dir.glob("*"):
|
||
if file_path.is_file() and file_path.stat().st_mtime < cutoff_time:
|
||
try:
|
||
file_path.unlink()
|
||
deleted_count += 1
|
||
logger.info(f"删除旧文件: {file_path.name}")
|
||
except Exception as e:
|
||
logger.error(f"删除文件失败 {file_path.name}: {e}")
|
||
|
||
return deleted_count
|
||
|
||
def close(self):
|
||
"""关闭服务，释放资源"""
|
||
if self._sdk:
|
||
self._sdk.close()
|
||
self._sdk = None
|
||
logger.info("企业微信SDK已关闭")
|
||
|
||
def __enter__(self):
|
||
"""上下文管理器入口"""
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
"""上下文管理器退出，自动释放资源"""
|
||
self.close()
|
||
|
||
# 全局服务实例（单例模式）
|
||
_wework_service: Optional[WeworkService] = None
|
||
|
||
def get_wework_service() -> WeworkService:
|
||
"""获取企业微信服务实例（单例）"""
|
||
global _wework_service
|
||
if _wework_service is None:
|
||
_wework_service = WeworkService()
|
||
return _wework_service
|
||
|
||
def close_wework_service():
|
||
"""关闭企业微信服务"""
|
||
global _wework_service
|
||
if _wework_service:
|
||
_wework_service.close()
|
||
_wework_service = None
|
||
|
||
if __name__ == "__main__":
|
||
# 简单测试
|
||
import logging
|
||
logging.basicConfig(level=logging.INFO)
|
||
|
||
service = WeworkService()
|
||
print(f"企业微信服务初始化成功")
|
||
print(f"企业ID: {service.corpid}")
|
||
print(f"媒体目录: {service.media_dir}")
|
||
|
||
# 测试获取媒体文件列表
|
||
files = service.get_media_files()
|
||
print(f"现有媒体文件: {len(files)} 个")
|
||
|
||
service.close()
|