feat: 新增 decode-images 子命令(批量解密 .dat 图片到明文图片树)
## 问题 \`decode_image.py\` 目前只有 \`decrypt_dat_file()\` 单文件 API,以及 \`monitor_web\` 在收到新消息时\"按需解一张\"的路径。**没有\"一次性扫 attach 目录、产出明文图片树到固定路径\"的批量入口**。结果是任何想把微信图片做下游消费(数据分析、搜索索引、归档、第三方 viewer)的用户都得各自写一遍 walk + decrypt 的 wrapper,且各自约定输出布局,生态不收敛。 ## 修复 - \`decode_image.py\` 新增 \`decode_all_dats(attach_dir, out_dir, aes_key, xor_key, force, on_file)\` 函数,扫描 \`<attach_dir>/<chat_hash>/<YYYY-MM>/Img/*.dat\` 并镜像产出 \`<out_dir>/<chat_hash>/<YYYY-MM>/<file_md5>.<ext>\`。 - \`main.py\` 新增 \`decode-images\` 子命令(早路由,跳过 \`check_wechat_running\` 和 \`ensure_keys\` —— 这条路径只读 \`.dat\` 文件,既不需要微信进程也不需要 DB 密钥)。 设计选择: - **输出布局 1:1 镜像 attach**,只做最小 path massage(去 \`Img/\`、去 \`_t/_h\` 缩略图后缀、换扩展名),不发明新结构。下游能用 \`md5(username)\` 反推路径,无需读 mapping 文件。 - **幂等性 = 按 basename 存在性 skip**,不做 mtime 比较 —— \`.dat\` 是 content-hash 命名(\`file_md5 = 文件内容 md5\`),实际上 write-once。\`--force\` 强制重解。 - **原子写**:解密先写 \`<basename>.<ext>.tmp\`(同目录),\`os.replace\` 到正式路径。中断不留半个 jpg。残留 \`.tmp\` 不会被 skip 误判(glob 显式排除)。 - **错误隔离**:单文件失败计入 \`failed\` 继续下一个,stderr 打 \`[WARN]\` 指出相对路径。退出码 2 表示\"部分失败,产物部分可用\"。 - **V2 无 key**:计入 \`skipped_no_key\` 而非 \`failed\` —— 这是可恢复状态(跑 \`find_image_key_macos.py\` / \`find_image_key.py\` 后重跑即可),跟\"真失败\"区分对待。V1 / 老 XOR 不依赖 \`image_aes_key\`。 - **wxgf 容器**只产 \`.hevc\` 裸流,**不**做 mp4 转换:上游不引入 ffmpeg subprocess 依赖,转换是消费层职责。 - **CLI override**:\`--attach-dir\` / \`--decoded-dir\` / \`--aes-key\` / \`--xor-key\` / \`--force\` 都可覆盖 \`config.json\`,适合 CI / 多账号 / 容器化场景。 ## 测试 新文件 \`tests/test_decode_images_batch.py\`,13 个新测试: - \`PathParsingTests\` (4):glob 命中 / \`_t\` 后缀剥离 / \`_h\` 后缀剥离 / chat_hash + YYYY-MM 镜像 - \`IdempotentTests\` (3):已存在跳过 / \`--force\` 覆写 / 残留 \`.tmp\` 不误判 - \`AtomicWriteTests\` (3):成功路径无 \`.tmp\` / decrypt 返回 None 无 \`.tmp\` / decrypt 抛异常无 \`.tmp\` - \`V2NoKeyTests\` (2):V2 + 无 key → skipped_no_key / V1 + 无 key 仍解码 - \`CallbackTests\` (1):\`on_file\` 回调每文件触发 基线 183 → 196 通过(+13 新增),0 回归。\`decrypt_dat_file\` 用 mock 隔离(避免依赖真实加密图片);\`is_v2_format\` 走真实 magic 检测路径。 ## 范围 - \`decode_image.py\`:新增 \`decode_all_dats\` 函数,134 行,纯加,不改任何现有 API。 - \`main.py\`:新增 \`_run_decode_images\` helper + 早路由 + 用法 hint,104 行加 2 行删。无 backward-compat 影响。 - \`tests/test_decode_images_batch.py\`:新增,295 行。合成 fixture(假 V1/V2 magic + mock decrypt_dat_file),不依赖真实加密素材。
This commit is contained in:
Belugary
295
tests/test_decode_images_batch.py
Normal file
295
tests/test_decode_images_batch.py
Normal file
@@ -0,0 +1,295 @@
|
||||
"""decode_image.decode_all_dats() batch CLI 行为测试。
|
||||
|
||||
覆盖:
|
||||
- 路径扫描:glob 命中 attach/<chat_hash>/<YYYY-MM>/Img/*.dat
|
||||
- 路径解析:chat_hash / YYYY-MM 提取,_t / _h 后缀移除归并到原图 basename
|
||||
- 幂等性:目标 basename 已存在(任何扩展名)时跳过;--force 强制重解
|
||||
- 原子写:写到 tmp 再 os.replace;失败/异常路径不留 .tmp
|
||||
- V2 无 key:计入 skipped_no_key 而非 failed
|
||||
- 错误隔离:单文件异常不阻塞批次;返回失败计数
|
||||
|
||||
decrypt_dat_file 用 mock 隔离(避免依赖真实加密图片);is_v2_format
|
||||
单独覆盖真实 magic 检测路径。
|
||||
"""
|
||||
import os
|
||||
import struct
|
||||
import tempfile
|
||||
import unittest
|
||||
from contextlib import redirect_stderr
|
||||
import io
|
||||
from unittest.mock import patch
|
||||
|
||||
import decode_image
|
||||
|
||||
|
||||
def _write(path, data):
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "wb") as f:
|
||||
f.write(data)
|
||||
|
||||
|
||||
def _v2_magic_bytes():
|
||||
# 仅用于让 is_v2_format() 返回 True
|
||||
return decode_image.V2_MAGIC_FULL + struct.pack("<LL", 0, 0) + b"\x00"
|
||||
|
||||
|
||||
def _v1_magic_bytes():
|
||||
return decode_image.V1_MAGIC_FULL + struct.pack("<LL", 0, 0) + b"\x00"
|
||||
|
||||
|
||||
class _MockedDecrypt:
|
||||
"""mock decrypt_dat_file:不真解密,只往 tmp 写一个 marker 字节串然后返回 (tmp, ext)。
|
||||
|
||||
通过实例化时配置返回的 ext / 是否抛异常 / 是否返回 (None, None),覆盖
|
||||
各种成功/失败路径。
|
||||
"""
|
||||
def __init__(self, ext="jpg", marker=b"DECODED", returns_none=False, raises=None):
|
||||
self.ext = ext
|
||||
self.marker = marker
|
||||
self.returns_none = returns_none
|
||||
self.raises = raises
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, dat_path, out_path=None, aes_key=None, xor_key=0x88):
|
||||
self.calls.append((dat_path, out_path, aes_key, xor_key))
|
||||
if self.raises:
|
||||
raise self.raises
|
||||
if self.returns_none:
|
||||
return None, None
|
||||
# 写 tmp(decode_all_dats 期望我们写完才能 os.replace)
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(self.marker)
|
||||
return out_path, self.ext
|
||||
|
||||
|
||||
def _make_dat(attach_dir, chat_hash, ym, basename, content=None):
|
||||
"""在 attach_dir 下造一个 .dat 文件,返回完整路径。content 默认是非 V2 占位。"""
|
||||
if content is None:
|
||||
content = b"\x00\x00\x00\x00" # 非 V2 / 非 V1 magic
|
||||
p = os.path.join(attach_dir, chat_hash, ym, "Img", f"{basename}.dat")
|
||||
_write(p, content)
|
||||
return p
|
||||
|
||||
|
||||
class PathParsingTests(unittest.TestCase):
|
||||
"""路径扫描 / 解析 / _t _h 归并。"""
|
||||
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
self.attach = os.path.join(self._tmp.name, "attach")
|
||||
self.out = os.path.join(self._tmp.name, "out")
|
||||
|
||||
def test_finds_dat_files_under_chat_month_img(self):
|
||||
_make_dat(self.attach, "hash1", "2026-01", "abc123")
|
||||
_make_dat(self.attach, "hash2", "2026-02", "def456")
|
||||
mock = _MockedDecrypt()
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["total"], 2)
|
||||
self.assertEqual(stats["decoded"], 2)
|
||||
self.assertEqual(stats["failed"], 0)
|
||||
|
||||
def test_strips_t_suffix(self):
|
||||
_make_dat(self.attach, "hash1", "2026-01", "abc123_t")
|
||||
mock = _MockedDecrypt(ext="png")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
# 期望产出: out/hash1/2026-01/abc123.png(_t 已被剥)
|
||||
produced = os.path.join(self.out, "hash1", "2026-01", "abc123.png")
|
||||
self.assertTrue(os.path.exists(produced), f"missing: {produced}")
|
||||
|
||||
def test_strips_h_suffix(self):
|
||||
_make_dat(self.attach, "hash1", "2026-01", "abc_h")
|
||||
mock = _MockedDecrypt(ext="jpg")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
produced = os.path.join(self.out, "hash1", "2026-01", "abc.jpg")
|
||||
self.assertTrue(os.path.exists(produced))
|
||||
|
||||
def test_mirrors_chat_and_month(self):
|
||||
_make_dat(self.attach, "abcdef0123456789", "2026-04", "img1")
|
||||
mock = _MockedDecrypt(ext="jpg")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
produced = os.path.join(self.out, "abcdef0123456789", "2026-04", "img1.jpg")
|
||||
self.assertTrue(os.path.exists(produced))
|
||||
|
||||
|
||||
class IdempotentTests(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
self.attach = os.path.join(self._tmp.name, "attach")
|
||||
self.out = os.path.join(self._tmp.name, "out")
|
||||
|
||||
def test_existing_target_basename_skipped(self):
|
||||
_make_dat(self.attach, "hash1", "2026-01", "img1")
|
||||
# 预先放一个目标(任何扩展名)
|
||||
existing = os.path.join(self.out, "hash1", "2026-01", "img1.png")
|
||||
_write(existing, b"OLD")
|
||||
mock = _MockedDecrypt()
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["skipped"], 1)
|
||||
self.assertEqual(stats["decoded"], 0)
|
||||
self.assertEqual(len(mock.calls), 0, "decrypt_dat_file 不该被调用")
|
||||
# 目标内容未被改写
|
||||
with open(existing, "rb") as f:
|
||||
self.assertEqual(f.read(), b"OLD")
|
||||
|
||||
def test_force_overrides_skip(self):
|
||||
_make_dat(self.attach, "hash1", "2026-01", "img1")
|
||||
existing = os.path.join(self.out, "hash1", "2026-01", "img1.png")
|
||||
_write(existing, b"OLD")
|
||||
mock = _MockedDecrypt(ext="jpg", marker=b"NEW")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16,
|
||||
force=True, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["decoded"], 1)
|
||||
self.assertEqual(stats["skipped"], 0)
|
||||
# 新文件以新 ext 落盘
|
||||
new_file = os.path.join(self.out, "hash1", "2026-01", "img1.jpg")
|
||||
self.assertTrue(os.path.exists(new_file))
|
||||
|
||||
def test_skip_ignores_tmp_files(self):
|
||||
"""残留的 .tmp 不应该被当成"已存在目标"误判跳过。"""
|
||||
_make_dat(self.attach, "hash1", "2026-01", "img1")
|
||||
# 模拟之前一次中断留下的 .tmp
|
||||
leftover_tmp = os.path.join(self.out, "hash1", "2026-01", "img1.unknown.tmp")
|
||||
_write(leftover_tmp, b"PARTIAL")
|
||||
mock = _MockedDecrypt(ext="jpg")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["decoded"], 1, "残留 .tmp 不应该阻止重解")
|
||||
|
||||
|
||||
class AtomicWriteTests(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
self.attach = os.path.join(self._tmp.name, "attach")
|
||||
self.out = os.path.join(self._tmp.name, "out")
|
||||
_make_dat(self.attach, "hash1", "2026-01", "img1")
|
||||
|
||||
def test_success_path_no_tmp_leftover(self):
|
||||
mock = _MockedDecrypt(ext="jpg")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
target_dir = os.path.join(self.out, "hash1", "2026-01")
|
||||
leftovers = [f for f in os.listdir(target_dir) if f.endswith(".tmp")]
|
||||
self.assertEqual(leftovers, [], "成功路径不应有 .tmp 残留")
|
||||
|
||||
def test_decrypt_returns_none_no_tmp_leftover(self):
|
||||
mock = _MockedDecrypt(returns_none=True)
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["failed"], 1)
|
||||
# decrypt 没写 tmp(returns_none=True 时也不写),所以目录可能不存在或为空
|
||||
target_dir = os.path.join(self.out, "hash1", "2026-01")
|
||||
if os.path.isdir(target_dir):
|
||||
leftovers = [f for f in os.listdir(target_dir) if f.endswith(".tmp")]
|
||||
self.assertEqual(leftovers, [])
|
||||
|
||||
def test_decrypt_raises_no_tmp_leftover(self):
|
||||
mock = _MockedDecrypt(raises=RuntimeError("synthetic decrypt failure"))
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key="x" * 16, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["failed"], 1)
|
||||
target_dir = os.path.join(self.out, "hash1", "2026-01")
|
||||
if os.path.isdir(target_dir):
|
||||
leftovers = [f for f in os.listdir(target_dir) if f.endswith(".tmp")]
|
||||
self.assertEqual(leftovers, [], "异常路径必须清理 .tmp")
|
||||
|
||||
|
||||
class V2NoKeyTests(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self._tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self._tmp.cleanup)
|
||||
self.attach = os.path.join(self._tmp.name, "attach")
|
||||
self.out = os.path.join(self._tmp.name, "out")
|
||||
|
||||
def test_v2_dat_with_no_aes_key_skipped(self):
|
||||
_make_dat(self.attach, "hash1", "2026-01", "v2img", content=_v2_magic_bytes())
|
||||
mock = _MockedDecrypt()
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key=None, progress_every=None,
|
||||
)
|
||||
self.assertEqual(stats["skipped_no_key"], 1)
|
||||
self.assertEqual(stats["decoded"], 0)
|
||||
self.assertEqual(stats["failed"], 0)
|
||||
self.assertEqual(len(mock.calls), 0, "无 key 的 V2 文件不应该走 decrypt_dat_file")
|
||||
|
||||
def test_v1_dat_with_no_aes_key_still_decoded(self):
|
||||
"""V1 用固定 AES key,不需要 image_aes_key,仍应被处理。"""
|
||||
_make_dat(self.attach, "hash1", "2026-01", "v1img", content=_v1_magic_bytes())
|
||||
mock = _MockedDecrypt(ext="jpg")
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
stats = decode_image.decode_all_dats(
|
||||
self.attach, self.out, aes_key=None, progress_every=None,
|
||||
)
|
||||
# is_v2_format 只识别 V2(纯 V2 magic),V1 不算 V2,所以会进入 decrypt 流程
|
||||
self.assertEqual(stats["decoded"], 1)
|
||||
self.assertEqual(stats["skipped_no_key"], 0)
|
||||
|
||||
|
||||
class CallbackTests(unittest.TestCase):
|
||||
|
||||
def test_on_file_callback_fires_per_file(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
attach = os.path.join(tmp, "attach")
|
||||
out = os.path.join(tmp, "out")
|
||||
_make_dat(attach, "h1", "2026-01", "a")
|
||||
_make_dat(attach, "h1", "2026-01", "b")
|
||||
_make_dat(attach, "h1", "2026-01", "c")
|
||||
events = []
|
||||
mock = _MockedDecrypt()
|
||||
with patch.object(decode_image, "decrypt_dat_file", mock), \
|
||||
redirect_stderr(io.StringIO()):
|
||||
decode_image.decode_all_dats(
|
||||
attach, out, aes_key="x" * 16, progress_every=None,
|
||||
on_file=lambda i, total, p, status, fmt: events.append(status),
|
||||
)
|
||||
self.assertEqual(len(events), 3)
|
||||
self.assertTrue(all(s == "decoded" for s in events))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user