diff --git a/db/message_storage.py b/db/message_storage.py index 9cd8723..b72eaa4 100644 --- a/db/message_storage.py +++ b/db/message_storage.py @@ -209,6 +209,20 @@ class MessageStorageDB(BaseDBOperator): print(f"更新消息图片路径出错: {e}") return False + def update_message_image_file_path(self, message_id, image_path): + try: + sql = """ + UPDATE messages + SET image_path = %s + WHERE message_id = %s + """ + params = (image_path, message_id) + result = self.execute_update(sql, params) + return result + except Exception as e: + print(f"更新消息图片文件路径出错: {e}") + return False + def get_hourly_message_trend(self, group_id: str = None, days: int = 1) -> List[Dict]: """获取指定群组的按小时消息趋势数据 diff --git a/utils/wechat/message_to_db.py b/utils/wechat/message_to_db.py index ba73e96..0e91a14 100644 --- a/utils/wechat/message_to_db.py +++ b/utils/wechat/message_to_db.py @@ -4,6 +4,8 @@ from datetime import datetime, timedelta import xml.etree.ElementTree as ET import concurrent.futures # 添加线程池支持 import os +import base64 +import imghdr from db.connection import DBConnectionManager from db.levels_db import LevelsDBOperator @@ -125,13 +127,38 @@ class MessageStorage: cdnthumburl = img_elem.attrib.get("cdnthumburl", "") base64_str = asyncio.run(self.client.download_image(aeskey=aeskey, cdnmidimgurl=cdnthumburl)) if base64_str: - self.message_db.update_message_image_path(msg.msg_id, "data:image/png;base64," + base64_str) + group_dir = os.path.join(self.image_dir, msg.roomid or "unknown") + if not os.path.exists(group_dir): + os.makedirs(group_dir, exist_ok=True) + data = base64.b64decode(base64_str) + kind = imghdr.what(None, h=data) + ext = "jpg" if kind == "jpeg" else (kind if kind else "png") + file_name = f"{msg.msg_id}.{ext}" + file_path = os.path.join(group_dir, file_name) + skipped = False + if not os.path.exists(file_path): + try: + with open(file_path, "wb") as f: + f.write(data) + except Exception as e: + return { + 'success': False, + 'message_id': msg.msg_id, + 'roomid': msg.roomid, + 'sender': msg.sender, + 'error': f"图片保存失败: {e}" + } + else: + skipped = True + web_path = f"/static/images/{msg.roomid}/{file_name}" + self.message_db.update_message_image_file_path(msg.msg_id, web_path) return { 'success': True, 'message_id': msg.msg_id, 'roomid': msg.roomid, 'sender': msg.sender, - 'file_path': base64_str + 'file_path': web_path, + 'skipped': skipped } else: return {