From c29e8dd868de31443462358fe0bee9b3ca86fcbb Mon Sep 17 00:00:00 2001
From: Belugary <53219544+Belugary@users.noreply.github.com>
Date: Tue, 5 May 2026 17:05:53 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20ImageResolver=20=E6=94=AF=E6=8C=81?=
 =?UTF-8?q?=E5=BE=AE=E4=BF=A1=204.0+=20V2=20=E5=8A=A0=E5=AF=86=E5=9B=BE?=
 =?UTF-8?q?=E7=89=87=E6=A0=BC=E5=BC=8F=20(#61)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ImageResolver.decode_image 之前只调 xor_decrypt_file（老格式 XOR-only
路径），微信 4.0+（2025-08+）已经改用 V2 AES-128-ECB + XOR 混合加密，
导致 mcp_server.py 注册的 decode_image MCP 工具对 V2 .dat 文件返回的
"解密"内容是错的——Claude AI 通过 MCP 调用看不到 V2 时代的图片。

monitor_web.py 早已正确处理 V2（line 41-42, 791-795：从 _cfg 读
image_aes_key / image_xor_key 后调 decrypt_dat_file 自动 magic 分发），
本次把 MCP 路径补齐，行为与 monitor_web.py 对齐。

改动：
- ImageResolver.__init__ 增加 aes_key=None, xor_key=0x88 关键字参数
  （默认值保持向后兼容，老调用方无需改动）
- ImageResolver.decode_image 把 xor_decrypt_file 换成 decrypt_dat_file，
  按 magic 自动分发 V2 / V1 / 老 XOR
- V2 文件 + 缺 aes_key 时早期返回结构化错误信息，避免在 v2_decrypt_file
  内静默失败成笼统的"解密失败"
- v2_decrypt_file 入口接受 xor_key 字符串形式（int(_, 0) 解析），
  与 aes_key 已有的 str→bytes 处理对称，允许 config.json 写 "0x88"
- mcp_server.py 实例化时从 _cfg 读 image_aes_key / image_xor_key 注入

兼容性：
- ImageResolver 老调用方（不传 keys）继续走老 XOR 路径，零 breaking
- V1 magic（\x07\x08V1）不会被 is_v2_format 拦截，走 decrypt_dat_file
  内置固定 key，所以 aes_key=None 也能解 V1 文件
- 整 repo 只有 mcp_server.py 一处生产调用 ImageResolver(...)，已 grep 确认

测试覆盖（11 个新测试，tests/test_decode_image_v2.py）：
- v2_decrypt_file 合成数据 round-trip 字节级相等
- decrypt_dat_file 按 magic 自动分发 V2 / V1 / 老 XOR 三条路径
- aes_key 接受 str（来自 config.json）和 bytes 两种形式
- xor_key 接受 str（如 "0x88"）和 int 两种形式
- V2 wxgf 裸流返回 fmt='hevc'（HEVC→JPEG 转换是 monitor_web 职责，
  不在 ImageResolver 内做，保留 .hevc 输出）
- ImageResolver 端到端：from local_id to decrypted file
- ImageResolver(aes_key=None) + V1 文件走固定 key 路径
- ImageResolver(aes_key=None) + V2 文件返回 success=False + 友好错误
- ImageResolver 默认参数 + 老 XOR .dat 保持向后兼容

测试 46 个全部通过（11 新 + 35 旧）。

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 decode_image.py               |  22 ++-
 mcp_server.py                 |   7 +-
 tests/test_decode_image_v2.py | 280 ++++++++++++++++++++++++++++++++++
 3 files changed, 303 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_decode_image_v2.py

diff --git a/decode_image.py b/decode_image.py
index f9edbc4..870bca2 100644
--- a/decode_image.py
+++ b/decode_image.py
@@ -118,7 +118,7 @@ def v2_decrypt_file(dat_path, out_path=None, aes_key=None, xor_key=0x88):
         dat_path: V2 .dat 文件路径
         out_path: 输出路径 (None 则自动命名)
         aes_key: 16 字节 AES key (bytes 或 str)
-        xor_key: XOR key (int, 默认 0x88)
+        xor_key: XOR key (int 或可被 int(_, 0) 解析的 str, 默认 0x88)
 
     Returns:
         (output_path, format) 或 (None, None)
@@ -135,6 +135,10 @@ def v2_decrypt_file(dat_path, out_path=None, aes_key=None, xor_key=0x88):
     if len(aes_key) < 16:
         return None, None
 
+    # 与 aes_key 的 str→bytes 处理对称: 允许 config.json 写 "0x88" / "136" 等字符串形式
+    if isinstance(xor_key, str):
+        xor_key = int(xor_key, 0)
+
     with open(dat_path, 'rb') as f:
         data = f.read()
 
@@ -299,17 +303,21 @@ def extract_md5_from_packed_info(blob):
 class ImageResolver:
     """封装从 local_id 到图片文件的完整解析链"""
 
-    def __init__(self, wechat_base_dir, decoded_image_dir, cache):
+    def __init__(self, wechat_base_dir, decoded_image_dir, cache, aes_key=None, xor_key=0x88):
         """
         Args:
             wechat_base_dir: 微信数据根目录 (如 D:\\xwechat_files\\<wxid>)
             decoded_image_dir: 解密图片输出目录
             cache: DBCache 实例，用于解密 message_resource.db
+            aes_key: V2 格式的 AES key (16 字节 str/bytes)，None 表示不支持 V2 文件
+            xor_key: XOR key (int, 默认 0x88)，用于 V2 文件的 XOR 段
         """
         self.base_dir = wechat_base_dir
         self.attach_dir = os.path.join(wechat_base_dir, "msg", "attach")
         self.out_dir = decoded_image_dir
         self.cache = cache
+        self.aes_key = aes_key
+        self.xor_key = xor_key
 
     def get_image_md5(self, local_id):
         """通过 local_id 查 message_resource.db 获取图片文件 MD5"""
@@ -379,13 +387,17 @@ class ImageResolver:
                 selected = f
                 break
 
-        # 3. 解密
+        # 3. 解密 (decrypt_dat_file 会按 magic 自动分发 V2 / V1 / 老 XOR)
         out_name = f"{file_md5}"
         out_path_base = os.path.join(self.out_dir, out_name)
 
-        result_path, fmt = xor_decrypt_file(selected, f"{out_path_base}.tmp")
+        # 提前拦截以给出具体错误信息;否则会在 v2_decrypt_file 内 silent-fail 成笼统的"解密失败"
+        if is_v2_format(selected) and not self.aes_key:
+            return {'success': False, 'error': f'V2 格式 .dat 文件需要 AES key (文件: {selected})', 'md5': file_md5}
+
+        result_path, fmt = decrypt_dat_file(selected, f"{out_path_base}.tmp", self.aes_key, self.xor_key)
         if not result_path:
-            return {'success': False, 'error': f'无法检测 XOR key (文件: {selected})', 'md5': file_md5}
+            return {'success': False, 'error': f'解密失败 (文件: {selected})', 'md5': file_md5}
 
         # 重命名为正确扩展名
         final_path = f"{out_path_base}.{fmt}"
diff --git a/mcp_server.py b/mcp_server.py
index 71e75a8..901dd5b 100644
--- a/mcp_server.py
+++ b/mcp_server.py
@@ -1664,7 +1664,12 @@ def get_new_messages() -> str:
 
 # ============ 图片解密 ============
 
-_image_resolver = ImageResolver(WECHAT_BASE_DIR, DECODED_IMAGE_DIR, _cache)
+_image_aes_key = _cfg.get("image_aes_key")  # V2 格式 AES key (从微信内存提取)
+_image_xor_key = _cfg.get("image_xor_key", 0x88)
+_image_resolver = ImageResolver(
+    WECHAT_BASE_DIR, DECODED_IMAGE_DIR, _cache,
+    aes_key=_image_aes_key, xor_key=_image_xor_key,
+)
 
 
 @mcp.tool()
diff --git a/tests/test_decode_image_v2.py b/tests/test_decode_image_v2.py
new file mode 100644
index 0000000..2c4112d
--- /dev/null
+++ b/tests/test_decode_image_v2.py
@@ -0,0 +1,280 @@
+"""ImageResolver 在 V2 加密格式下的端到端解密测试。
+
+覆盖:
+- v2_decrypt_file 能正确还原 AES-ECB + XOR 混合加密的合成数据
+- decrypt_dat_file 按 magic 自动分发 V2 / V1 / 老 XOR 三条路径
+- ImageResolver 通过 __init__ 注入 aes_key/xor_key 后,能端到端解密 V2 .dat
+- 没传 aes_key 时遇到 V2 文件返回结构化错误,而不是 crash 或返回错误数据
+- 默认参数下老 XOR 路径不受影响,保持向后兼容
+"""
+import hashlib
+import os
+import sqlite3
+import struct
+import tempfile
+import unittest
+
+from Crypto.Cipher import AES
+from Crypto.Util import Padding
+
+from decode_image import (
+    V1_MAGIC_FULL,
+    V2_MAGIC_FULL,
+    ImageResolver,
+    decrypt_dat_file,
+    v2_decrypt_file,
+)
+
+
+# 测试用 16 字节 AES key (任意值,仅用于合成测试数据)
+TEST_AES_KEY = b'1234567890abcdef'
+TEST_XOR_KEY = 0x37
+# 最小可识别的 PNG payload (含 IHDR 和 IEND chunk),长度 88 字节
+TEST_PNG_PAYLOAD = (
+    b'\x89PNG\r\n\x1a\n'
+    + b'\x00\x00\x00\rIHDR'
+    + b'\x00' * 64
+    + b'IEND\xaeB`\x82'
+)
+
+
+def _build_v2_dat(plaintext, aes_size, xor_size,
+                  aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY,
+                  magic=V2_MAGIC_FULL):
+    """构造合成的 V2 / V1 .dat 字节串。
+
+    布局: [6B magic][4B aes_size LE][4B xor_size LE][1B pad][AES-ECB][raw][XOR]
+    aes_size / xor_size 是明文字段长度,AES 段做 PKCS7 padding 后向上对齐到 16 倍数。
+    """
+    if aes_size + xor_size > len(plaintext):
+        raise ValueError("aes_size + xor_size 超过 plaintext 长度")
+    aes_plain = plaintext[:aes_size]
+    raw_plain = plaintext[aes_size:len(plaintext) - xor_size]
+    xor_plain = plaintext[len(plaintext) - xor_size:]
+
+    cipher = AES.new(aes_key[:16], AES.MODE_ECB)
+    aes_cipher = cipher.encrypt(Padding.pad(aes_plain, AES.block_size))
+    xor_cipher = bytes(b ^ xor_key for b in xor_plain)
+
+    header = magic + struct.pack('<LL', aes_size, xor_size) + b'\x00'
+    return header + aes_cipher + raw_plain + xor_cipher
+
+
+class _FakeCache:
+    """ImageResolver 测试用最小缓存桩,绕过真实 DB 解密。"""
+
+    def __init__(self, mapping):
+        self._mapping = mapping
+
+    def get(self, rel_key):
+        return self._mapping.get(rel_key)
+
+
+def _make_resource_db(path, local_id, file_md5):
+    """构造最小 message_resource.db,只含一条 packed_info 记录。
+
+    packed_info 里嵌入 extract_md5_from_packed_info 期望的 protobuf marker
+    (\\x12\\x22\\x0a\\x20) 加 32 字节 ASCII hex MD5。
+    """
+    marker = b'\x12\x22\x0a\x20'
+    packed = b'\x00' * 8 + marker + file_md5.encode('ascii') + b'\x00' * 4
+    conn = sqlite3.connect(path)
+    try:
+        conn.execute(
+            "CREATE TABLE MessageResourceInfo (local_id INTEGER PRIMARY KEY, packed_info BLOB)"
+        )
+        conn.execute(
+            "INSERT INTO MessageResourceInfo VALUES (?, ?)",
+            (local_id, packed),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+class TestV2DecryptSynthetic(unittest.TestCase):
+    """v2_decrypt_file / decrypt_dat_file 在合成数据上的正确性"""
+
+    def test_v2_round_trip_recovers_payload(self):
+        # 合成 V2 .dat 解密后字节级等于原始 payload
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(_build_v2_dat(TEST_PNG_PAYLOAD, aes_size=32, xor_size=16))
+
+            out_path, fmt = v2_decrypt_file(
+                dat_path, aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY
+            )
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'png')
+            with open(out_path, 'rb') as f:
+                self.assertEqual(f.read(), TEST_PNG_PAYLOAD)
+
+    def test_decrypt_dat_file_routes_v2_by_magic(self):
+        # decrypt_dat_file 看到 V2 magic 应自动走 V2 路径
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(_build_v2_dat(TEST_PNG_PAYLOAD, aes_size=32, xor_size=16))
+
+            out_path, fmt = decrypt_dat_file(
+                dat_path, aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY
+            )
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'png')
+
+    def test_decrypt_dat_file_v1_uses_fixed_key(self):
+        # V1 magic 走固定 key,即便外部不传 aes_key 也能解密
+        v1_fixed_key = b'cfcd208495d565ef'  # md5("0")[:16],由 v2_decrypt_file 内部使用
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(_build_v2_dat(
+                    TEST_PNG_PAYLOAD, aes_size=32, xor_size=16,
+                    aes_key=v1_fixed_key, magic=V1_MAGIC_FULL,
+                ))
+
+            out_path, fmt = decrypt_dat_file(
+                dat_path, aes_key=None, xor_key=TEST_XOR_KEY
+            )
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'png')
+
+    def test_decrypt_dat_file_legacy_xor_route(self):
+        # 老 XOR 格式 (无 V1/V2 magic),decrypt_dat_file 应回退到 xor_decrypt_file 不需要 aes_key
+        xor_key = 0x37
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(bytes(b ^ xor_key for b in TEST_PNG_PAYLOAD))
+
+            out_path, fmt = decrypt_dat_file(dat_path, aes_key=None)
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'png')
+
+    def test_v2_accepts_str_aes_key_from_config(self):
+        # 真实场景下 aes_key 来自 config.json,是 ASCII string 不是 bytes;
+        # v2_decrypt_file 内部应自行 encode,避免 TypeError
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(_build_v2_dat(TEST_PNG_PAYLOAD, aes_size=32, xor_size=16))
+
+            out_path, fmt = decrypt_dat_file(
+                dat_path, aes_key=TEST_AES_KEY.decode('ascii'), xor_key=TEST_XOR_KEY,
+            )
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'png')
+
+    def test_v2_accepts_str_xor_key_from_config(self):
+        # 与 aes_key 的 str 处理对称: config.json 里把 xor_key 写成 "0x88" / "136" 也应能正常解密
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(_build_v2_dat(TEST_PNG_PAYLOAD, aes_size=32, xor_size=16))
+
+            out_path, fmt = decrypt_dat_file(
+                dat_path, aes_key=TEST_AES_KEY, xor_key=hex(TEST_XOR_KEY),
+            )
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'png')
+
+    def test_v2_wxgf_payload_returns_hevc_format(self):
+        # 微信 V2 动图 (wxgf 裸流 HEVC) 解密后 fmt='hevc',输出文件以 .hevc 结尾;
+        # 当前 ImageResolver 不再向 JPEG 转 (那是 monitor_web 的职责),保持原样输出。
+        wxgf_payload = b'wxgf' + b'\x00' * 84  # 88 字节,与 PNG payload 同长度,避免改 aes/xor sizes
+        with tempfile.TemporaryDirectory() as td:
+            dat_path = os.path.join(td, "test.dat")
+            with open(dat_path, 'wb') as f:
+                f.write(_build_v2_dat(wxgf_payload, aes_size=32, xor_size=16))
+
+            out_path, fmt = decrypt_dat_file(
+                dat_path, aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY,
+            )
+            self.assertIsNotNone(out_path)
+            self.assertEqual(fmt, 'hevc')
+            self.assertTrue(out_path.endswith('.hevc'))
+
+
+class TestImageResolverV2(unittest.TestCase):
+    """ImageResolver 端到端:从 local_id 到解密文件,验证 V2 keys 注入路径"""
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self._tmp.cleanup)
+        tmp = self._tmp.name
+
+        self.wechat_base = os.path.join(tmp, "wechat")
+        self.out_dir = os.path.join(tmp, "decoded")
+        os.makedirs(self.out_dir, exist_ok=True)
+
+        self.username = "wxid_test123"
+        self.local_id = 42
+        self.file_md5 = "0123456789abcdef0123456789abcdef"
+
+        username_hash = hashlib.md5(self.username.encode()).hexdigest()
+        img_dir = os.path.join(
+            self.wechat_base, "msg", "attach", username_hash, "2025-08", "Img"
+        )
+        os.makedirs(img_dir, exist_ok=True)
+        self.dat_path = os.path.join(img_dir, f"{self.file_md5}.dat")
+        with open(self.dat_path, 'wb') as f:
+            f.write(_build_v2_dat(TEST_PNG_PAYLOAD, aes_size=32, xor_size=16))
+
+        self.db_path = os.path.join(tmp, "message_resource.db")
+        _make_resource_db(self.db_path, self.local_id, self.file_md5)
+        self.cache = _FakeCache({"message/message_resource.db": self.db_path})
+
+    def test_decode_image_v2_with_keys(self):
+        resolver = ImageResolver(
+            self.wechat_base, self.out_dir, self.cache,
+            aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY,
+        )
+        result = resolver.decode_image(self.username, self.local_id)
+        self.assertTrue(result['success'], msg=result)
+        self.assertEqual(result['format'], 'png')
+        self.assertEqual(result['md5'], self.file_md5)
+        with open(result['path'], 'rb') as f:
+            self.assertEqual(f.read(), TEST_PNG_PAYLOAD)
+
+    def test_decode_image_v2_missing_aes_key_returns_error(self):
+        # 没传 aes_key 时遇到 V2 文件应返回 success=False,而不是 crash 或写入错误文件
+        resolver = ImageResolver(
+            self.wechat_base, self.out_dir, self.cache, aes_key=None,
+        )
+        result = resolver.decode_image(self.username, self.local_id)
+        self.assertFalse(result['success'])
+        self.assertIn('AES key', result['error'])
+        self.assertEqual(result['md5'], self.file_md5)
+
+    def test_decode_image_default_args_preserve_legacy_xor(self):
+        # 默认参数 (aes_key=None) + 老 XOR .dat 应保持向后兼容
+        os.unlink(self.dat_path)
+        legacy_xor_key = 0x37
+        with open(self.dat_path, 'wb') as f:
+            f.write(bytes(b ^ legacy_xor_key for b in TEST_PNG_PAYLOAD))
+
+        resolver = ImageResolver(self.wechat_base, self.out_dir, self.cache)
+        result = resolver.decode_image(self.username, self.local_id)
+        self.assertTrue(result['success'], msg=result)
+        self.assertEqual(result['format'], 'png')
+
+    def test_decode_image_v1_no_aes_key_uses_fixed_key(self):
+        # V1 magic 不会被 is_v2_format guard 拦截 (V1 magic 是 \x07\x08V1, V2 是 \x07\x08V2);
+        # 即便 ImageResolver(aes_key=None), V1 文件也应通过 decrypt_dat_file 内置固定 key 解密
+        os.unlink(self.dat_path)
+        v1_fixed_key = b'cfcd208495d565ef'
+        with open(self.dat_path, 'wb') as f:
+            f.write(_build_v2_dat(
+                TEST_PNG_PAYLOAD, aes_size=32, xor_size=16,
+                aes_key=v1_fixed_key, magic=V1_MAGIC_FULL,
+            ))
+
+        resolver = ImageResolver(self.wechat_base, self.out_dir, self.cache, aes_key=None)
+        result = resolver.decode_image(self.username, self.local_id)
+        self.assertTrue(result['success'], msg=result)
+        self.assertEqual(result['format'], 'png')
+
+
+if __name__ == '__main__':
+    unittest.main()