#!/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()