From fa273b810dcd6ab14ecdd3516641cf2d7db5188b Mon Sep 17 00:00:00 2001 From: Banghao Chi Date: Sun, 8 Mar 2026 15:30:10 -0500 Subject: [PATCH] Improve chat history formatting --- mcp_server.py | 292 +++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 229 insertions(+), 63 deletions(-) diff --git a/mcp_server.py b/mcp_server.py index c6b4df0..c801d31 100644 --- a/mcp_server.py +++ b/mcp_server.py @@ -5,14 +5,15 @@ Based on FastMCP (stdio transport), reuses existing decryption. Runs on Windows Python (needs access to D:\ WeChat databases). """ -import os, sys, json, time, sqlite3, tempfile, struct, hashlib, atexit, re -import hmac as hmac_mod -from datetime import datetime -from Crypto.Cipher import AES -from mcp.server.fastmcp import FastMCP -import zstandard as zstd -from decode_image import ImageResolver -from key_utils import get_key_info, key_path_variants, strip_key_metadata +import os, sys, json, time, sqlite3, tempfile, struct, hashlib, atexit, re +import hmac as hmac_mod +from datetime import datetime +import xml.etree.ElementTree as ET +from Crypto.Cipher import AES +from mcp.server.fastmcp import FastMCP +import zstandard as zstd +from decode_image import ImageResolver +from key_utils import get_key_info, key_path_variants, strip_key_metadata # ============ 加密常量 ============ PAGE_SZ = 4096 @@ -216,8 +217,9 @@ atexit.register(_cache.cleanup) # ============ 联系人缓存 ============ -_contact_names = None # {username: display_name} -_contact_full = None # [{username, nick_name, remark}] +_contact_names = None # {username: display_name} +_contact_full = None # [{username, nick_name, remark}] +_self_username = None def _load_contacts_from(db_path): @@ -261,21 +263,32 @@ def get_contact_names(): return {} -def get_contact_full(): - global _contact_full - if _contact_full is None: - get_contact_names() - return _contact_full or [] +def get_contact_full(): + global _contact_full + if _contact_full is None: + get_contact_names() + return _contact_full or [] # ============ 辅助函数 ============ -def format_msg_type(t): - return { - 1: '文本', 3: '图片', 34: '语音', 42: '名片', - 43: '视频', 47: '表情', 48: '位置', 49: '链接/文件', - 50: '通话', 10000: '系统', 10002: '撤回', - }.get(t, f'type={t}') +def format_msg_type(t): + base_type, _ = _split_msg_type(t) + return { + 1: '文本', 3: '图片', 34: '语音', 42: '名片', + 43: '视频', 47: '表情', 48: '位置', 49: '链接/文件', + 50: '通话', 10000: '系统', 10002: '撤回', + }.get(base_type, f'type={t}') + + +def _split_msg_type(t): + try: + t = int(t) + except (TypeError, ValueError): + return 0, 0 + if t > 0xFFFFFFFF: + return t & 0xFFFFFFFF, t >> 32 + return t, 0 def resolve_username(chat_name): @@ -316,10 +329,10 @@ def _decompress_content(content, ct): return content -def _parse_message_content(content, local_type, is_group): - """解析消息内容，返回 (sender_id, text)""" - if content is None: - return '', '' +def _parse_message_content(content, local_type, is_group): + """解析消息内容，返回 (sender_id, text)""" + if content is None: + return '', '' if isinstance(content, bytes): return '', '(二进制内容)' @@ -328,7 +341,163 @@ def _parse_message_content(content, local_type, is_group): if is_group and ':\n' in content: sender, text = content.split(':\n', 1) - return sender, text + return sender, text + + +def _collapse_text(text): + if not text: + return '' + return re.sub(r'\s+', ' ', text).strip() + + +def _get_self_username(): + global _self_username + if _self_username is not None: + return _self_username + + names = get_contact_names() + account_dir = os.path.basename(os.path.dirname(DB_DIR)) + candidates = [account_dir] + + m = re.fullmatch(r'(.+)_([0-9a-fA-F]{4,})', account_dir) + if m: + candidates.insert(0, m.group(1)) + + for candidate in candidates: + if candidate and candidate in names: + _self_username = candidate + return _self_username + + _self_username = '' + return _self_username + + +def _load_name2id_maps(conn): + id_to_username = {} + username_to_id = {} + try: + rows = conn.execute("SELECT rowid, user_name FROM Name2Id").fetchall() + except sqlite3.Error: + return id_to_username, username_to_id + + for rowid, user_name in rows: + if not user_name: + continue + id_to_username[rowid] = user_name + username_to_id[user_name] = rowid + return id_to_username, username_to_id + + +def _display_name_for_username(username, names): + if not username: + return '' + if username == _get_self_username(): + return 'me' + return names.get(username, username) + + +def _resolve_sender_label(real_sender_id, sender_from_content, is_group, chat_username, chat_display_name, names, id_to_username): + sender_username = id_to_username.get(real_sender_id, '') + + if is_group: + if sender_username and sender_username != chat_username: + return _display_name_for_username(sender_username, names) + if sender_from_content: + return _display_name_for_username(sender_from_content, names) + return '' + + if sender_username == chat_username: + return chat_display_name + if sender_username: + return _display_name_for_username(sender_username, names) + return '' + + +def _resolve_quote_sender_label(ref_user, ref_display_name, is_group, chat_username, chat_display_name, names): + if is_group: + if ref_user: + return _display_name_for_username(ref_user, names) + return ref_display_name or '' + + if ref_user: + if ref_user == chat_username: + return chat_display_name + return 'me' + if ref_display_name: + if ref_display_name == chat_display_name: + return chat_display_name + return 'me' + return '' + + +def _format_app_message_text(content, local_type, is_group, chat_username, chat_display_name, names): + if not content or ' 160: + ref_content = ref_content[:160] + "..." + + quote_text = title or "[引用消息]" + if ref_content: + ref_label = _resolve_quote_sender_label( + ref_user, ref_display_name, is_group, chat_username, chat_display_name, names + ) + prefix = f"回复 {ref_label}: " if ref_label else "回复: " + quote_text += f"\n ↳ {prefix}{ref_content}" + return quote_text + + if app_type == 6: + return f"[文件] {title}" if title else "[文件]" + if app_type == 5: + return f"[链接] {title}" if title else "[链接]" + if app_type in (33, 36, 44): + return f"[小程序] {title}" if title else "[小程序]" + if title: + return f"[链接/文件] {title}" + return "[链接/文件]" + + +def _format_message_text(local_id, local_type, content, is_group, chat_username, chat_display_name, names): + sender_from_content, text = _parse_message_content(content, local_type, is_group) + base_type, _ = _split_msg_type(local_type) + + if base_type == 3: + text = f"[图片] (local_id={local_id})" + elif base_type == 47: + text = "[表情]" + elif base_type == 49: + text = _format_app_message_text( + text, local_type, is_group, chat_username, chat_display_name, names + ) or "[链接/文件]" + elif base_type != 1: + type_label = format_msg_type(local_type) + text = f"[{type_label}] {text}" if text else f"[{type_label}]" + + return sender_from_content, text # 消息 DB 的 rel_keys @@ -453,15 +622,16 @@ def get_chat_history(chat_name: str, limit: int = 50) -> str: if not db_path: return f"找不到 {display_name} 的消息记录（可能在未解密的DB中或无消息）" - conn = sqlite3.connect(db_path) - try: - rows = conn.execute(f""" - SELECT local_id, local_type, create_time, message_content, - WCDB_CT_message_content - FROM [{table_name}] - ORDER BY create_time DESC - LIMIT ? - """, (limit,)).fetchall() + conn = sqlite3.connect(db_path) + try: + id_to_username, _ = _load_name2id_maps(conn) + rows = conn.execute(f""" + SELECT local_id, local_type, create_time, real_sender_id, message_content, + WCDB_CT_message_content + FROM [{table_name}] + ORDER BY create_time DESC + LIMIT ? + """, (limit,)).fetchall() except Exception as e: conn.close() return f"查询失败: {e}" @@ -469,34 +639,30 @@ def get_chat_history(chat_name: str, limit: int = 50) -> str: if not rows: return f"{display_name} 无消息记录" - - lines = [] - for local_id, local_type, create_time, content, ct in reversed(rows): - time_str = datetime.fromtimestamp(create_time).strftime('%m-%d %H:%M') - - # zstd 解压 - content = _decompress_content(content, ct) - if content is None: - content = '(无法解压)' - - sender, text = _parse_message_content(content, local_type, is_group) - - if local_type == 3: - text = f"[图片] (local_id={local_id})" - elif local_type == 47: - text = "[表情]" - elif local_type != 1: - type_label = format_msg_type(local_type) - text = f"[{type_label}] {text}" if text else f"[{type_label}]" - - if text and len(text) > 500: - text = text[:500] + "..." - - if is_group and sender: - sender_name = names.get(sender, sender) - lines.append(f"[{time_str}] {sender_name}: {text}") - else: - lines.append(f"[{time_str}] {text}") + + lines = [] + for local_id, local_type, create_time, real_sender_id, content, ct in reversed(rows): + time_str = datetime.fromtimestamp(create_time).strftime('%m-%d %H:%M') + + # zstd 解压 + content = _decompress_content(content, ct) + if content is None: + content = '(无法解压)' + + sender, text = _format_message_text( + local_id, local_type, content, is_group, username, display_name, names + ) + + if text and len(text) > 500: + text = text[:500] + "..." + + sender_label = _resolve_sender_label( + real_sender_id, sender, is_group, username, display_name, names, id_to_username + ) + if sender_label: + lines.append(f"[{time_str}] {sender_label}: {text}") + else: + lines.append(f"[{time_str}] {text}") header = f"{display_name} 的最近 {len(lines)} 条消息" if is_group: