""" 消息服务 提供加载会话存档、解密、保存到数据库的整合功能 """ import json import os import sys import logging import base64 from pathlib import Path from datetime import datetime from typing import List, Dict, Any, Optional # 添加scripts目录到路径 sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts")) from wework_ctypes import WeWorkFinanceSDK, WeWorkSDKError from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import padding from cryptography.hazmat.backends import default_backend from app.models.message import Message logger = logging.getLogger(__name__) class MessageService: """消息服务，提供加载会话存档的整合功能""" 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", "") self.secret = secret or os.getenv("WEWORK_SECRET", "") # SDK实例（延迟初始化） self._sdk: Optional[WeWorkFinanceSDK] = None # RSA私钥路径 self.rsa_private_key_path = Path("cert/1_rsa.key") logger.info(f"消息服务初始化完成，企业ID: {self.corpid}") @property def sdk(self) -> WeWorkFinanceSDK: """获取SDK实例（延迟初始化）""" if self._sdk is None: try: # 动态库路径 lib_path = ( Path(__file__).parent.parent.parent / "lib" / "sdk" / "libWeWorkFinanceSdk_C.so" ) if not lib_path.exists(): raise FileNotFoundError(f"动态库不存在: {lib_path}") self._sdk = WeWorkFinanceSDK(self.corpid, self.secret, str(lib_path)) logger.info("企业微信SDK初始化成功") except Exception as e: logger.error(f"企业微信SDK初始化失败: {e}") raise return self._sdk def _load_rsa_private_key(self, 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(self, encrypted_data_base64: str, private_key: str) -> str: """ 使用RSA私钥解密Base64编码的数据 Args: encrypted_data_base64: Base64编码的加密数据 private_key: RSA私钥字符串 Returns: 解密后的字符串 """ try: # 加载私钥 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( self, encrypt_random_key: str, encrypt_chat_msg: str, publickey_ver: int = None ) -> dict: """ 解密聊天消息 Args: encrypt_random_key: 加密的随机密钥（Base64编码） encrypt_chat_msg: 加密的聊天消息 publickey_ver: publickey_ver版本号，用于动态匹配私钥 Returns: 解密后的消息字典 """ # 动态确定私钥路径 actual_key_path = self.rsa_private_key_path # 如果提供了publickey_ver，则动态匹配 if publickey_ver is not None: dynamic_key_path = f"cert/{publickey_ver}_rsa.key" dynamic_path = Path(dynamic_key_path) if dynamic_path.exists(): actual_key_path = dynamic_path # 1. 加载RSA私钥 private_key = self._load_rsa_private_key(str(actual_key_path)) # 2. 使用RSA私钥解密encrypt_random_key得到encrypt_key try: encrypt_key = self._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 = self.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} async def _decrypt_chat_data(self, chatdata: List[Dict[str, Any]]) -> tuple: """ 解密聊天数据 Args: chatdata: 原始聊天数据列表 Returns: (decrypted_chatdata, stats) decrypted_chatdata: 解密后的聊天数据列表 stats: 解密统计信息字典 """ decrypted_chatdata = [] stats = { "total_encrypted_messages": 0, "successful_decrypts": 0, "failed_decrypts": 0, } # 检查私钥文件是否存在 if not self.rsa_private_key_path.exists(): logger.warning(f"RSA私钥文件不存在: {self.rsa_private_key_path}") logger.warning("将跳过解密，保存原始加密数据") return chatdata, stats logger.info(f"使用RSA私钥: {self.rsa_private_key_path.absolute()}") logger.info("开始解密聊天数据...") 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 = self._decrypt_chat_message( encrypt_random_key, encrypt_chat_msg, 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: logger.info(f" 已解密 {stats['successful_decrypts']} 条消息...") except Exception as e: stats["failed_decrypts"] += 1 logger.warning(f"解密消息失败 (seq: {msg.get('seq', 'N/A')}): {e}") # 解密失败，跳过此消息（不添加到结果中） continue return decrypted_chatdata, stats async def _save_messages_to_database( self, 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: logger.info("没有消息需要存储到数据库") return stats logger.info(f"开始将 {len(messages)} 条消息存储到数据库...") try: # 批量处理消息 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实例 message_data = { "msgid": msg.get("msgid"), "action": msg.get("action"), "from_user": msg.get( "from" ), # 从原始数据的from字段读取，保存为from_user "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: logger.error(f"创建消息实例失败 (msgid: {msgid}): {e}") stats["failed"] += 1 # 批量插入新消息 if new_messages: try: await Message.insert_many(new_messages) stats["saved"] += len(new_messages) logger.info( f" 已存储 {len(new_messages)} 条消息 (进度: {i + len(batch)}/{len(messages)})" ) except Exception as e: logger.error(f"批量插入失败: {e}") stats["failed"] += len(new_messages) logger.info(f"数据库存储完成:") logger.info(f" 总计消息: {stats['total_messages']}") logger.info(f" 成功存储: {stats['saved']}") logger.info(f" 跳过重复: {stats['skipped_duplicates']}") logger.info(f" 存储失败: {stats['failed']}") except Exception as e: logger.error(f"数据库存储过程中发生错误: {e}") import traceback traceback.print_exc() stats["failed"] = len(messages) return stats async def get_messages( self, page: int = 1, page_size: int = 20, msgtype: Optional[str] = None, roomid: Optional[str] = None, from_user: Optional[str] = None, ) -> Dict[str, Any]: """ 查询消息列表，支持分页和过滤 Args: page: 页码（从1开始） page_size: 每页数量 msgtype: 消息类型过滤 roomid: 群聊ID过滤 from_user: 发送者过滤 Returns: 分页查询结果 """ try: # 构建查询条件 query = {} if msgtype: query["msgtype"] = msgtype if roomid: query["roomid"] = roomid if from_user: query["from_user"] = from_user # 使用from_user字段进行查询 # 计算跳过的文档数量 skip_count = (page - 1) * page_size # 执行查询 messages = await Message.find( query, skip=skip_count, limit=page_size, sort=[("msgtime", -1)], # 按消息时间倒序排列 ).to_list() # 获取总记录数 total_count = await Message.find(query).count() # 记录成功日志 logger.info( f"成功查询消息列表: 查询条件={query}, " f"页码={page}, 每页数量={page_size}, " f"总记录数={total_count}, 返回记录数={len(messages)}" ) # 准备响应数据 return { "success": True, "data": { "page": page, "page_size": page_size, "total": total_count, "total_pages": ( (total_count + page_size - 1) // page_size if page_size > 0 else 0 ), "messages": messages, }, } except Exception as e: logger.error(f"查询消息列表失败: {e}") import traceback traceback.print_exc() return { "success": False, "message": f"查询消息列表失败: {str(e)}", "data": { "page": page, "page_size": page_size, "total": 0, "total_pages": 0, "messages": [], }, } async def load_messages(self, seq: int = 0, limit: int = 100) -> Dict[str, Any]: """ 加载会话存档到数据库 完整流程：拉取 → 解密 → 保存 Args: seq: 起始seq值，首次使用0 limit: 拉取数量限制 Returns: 执行结果和统计信息 """ start_time = datetime.now() try: # 从环境变量获取代理配置 proxy = os.getenv("WEWORK_PROXY", "") passwd = os.getenv("WEWORK_PROXY_PASSWD", "") timeout = int(os.getenv("WEWORK_TIMEOUT", "10")) logger.info(f"开始加载会话存档: seq={seq}, limit={limit}") # 1. 拉取会话数据 logger.info("步骤1: 从企业微信拉取会话数据...") result_str = self.sdk.get_chat_data( seq=seq, limit=limit, proxy=proxy, passwd=passwd, timeout=timeout ) # 解析结果 result = json.loads(result_str) errcode = result.get("errcode", -1) errmsg = result.get("errmsg", "") if errcode != 0: raise Exception(f"拉取会话数据失败: errcode={errcode}, errmsg={errmsg}") chatdata = result.get("chatdata", []) total_fetched = len(chatdata) logger.info(f"成功拉取 {total_fetched} 条聊天记录") if total_fetched == 0: return { "success": True, "message": "没有新的聊天记录", "data": { "total_fetched": 0, "encrypted_count": 0, "decrypted_success": 0, "decrypted_failed": 0, "saved_count": 0, "duplicates_skipped": 0, "seq": seq, "max_seq": seq, "next_seq": seq, }, "execution_time": str(datetime.now() - start_time), } # 2. 解密消息 logger.info("步骤2: 解密加密消息...") decrypted_chatdata, decrypt_stats = await self._decrypt_chat_data(chatdata) # 3. 保存到数据库 logger.info("步骤3: 保存消息到数据库...") db_stats = await self._save_messages_to_database(decrypted_chatdata) # 计算最大seq值 max_seq = max( (msg.get("seq", 0) for msg in decrypted_chatdata), default=seq ) # 准备响应数据 execution_time = datetime.now() - start_time response_data = { "success": True, "message": "消息加载完成", "data": { "total_fetched": total_fetched, "encrypted_count": decrypt_stats["total_encrypted_messages"], "decrypted_success": decrypt_stats["successful_decrypts"], "decrypted_failed": decrypt_stats["failed_decrypts"], "saved_count": db_stats["saved"], "duplicates_skipped": db_stats["skipped_duplicates"], "failed_to_save": db_stats["failed"], "seq": seq, "max_seq": max_seq, "next_seq": max_seq + 1 if total_fetched > 0 else seq, }, "execution_time": str(execution_time), "timestamp": datetime.now().isoformat(), } logger.info(f"消息加载完成，耗时: {execution_time}") return response_data except json.JSONDecodeError as e: error_msg = f"解析JSON数据失败: {e}" logger.error(error_msg) return { "success": False, "message": error_msg, "execution_time": str(datetime.now() - start_time), } except Exception as e: error_msg = f"加载消息失败: {e}" logger.error(error_msg) import traceback traceback.print_exc() return { "success": False, "message": error_msg, "execution_time": str(datetime.now() - start_time), } def close(self): """关闭服务，释放资源""" if self._sdk: self._sdk.close() self._sdk = None logger.info("消息服务已关闭") # 全局服务实例（单例模式） _message_service: Optional[MessageService] = None def get_message_service() -> MessageService: """获取消息服务实例（单例）""" global _message_service if _message_service is None: _message_service = MessageService() return _message_service def close_message_service(): """关闭消息服务""" global _message_service if _message_service: _message_service.close() _message_service = None