From 216f44a99f6cbf471a8e3b9f772d351df44a403e Mon Sep 17 00:00:00 2001 From: Belugary <53219544+Belugary@users.noreply.github.com> Date: Tue, 12 May 2026 16:18:37 +0800 Subject: [PATCH] fix(image): reject corrupted V2 image when AES or XOR key is wrong (#81) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `v2_decrypt_file` previously wrote files to disk even when the keys were wrong, producing garbage output with no way for the caller to detect the failure: 1. Wrong AES key -> `detect_image_format` returns 'bin' (magic does not match any known format) -> a `.bin` file of random bytes was written. 2. Wrong XOR key with correct AES key -> file header looks like a valid jpg/png (the AES segment decrypts correctly) but the trailing XOR segment is scrambled -> callers get a half-valid image file that image viewers either render as truncated or fail to open. Both cases now return `(None, None)`: - `fmt == 'bin'` -> fail fast, no file written. - `xor_size >= 2` -> validate trailer magic by format: * jpg must end with FF D9 (EOI marker) * png must contain IEND chunk in the last 12 bytes Other formats (gif/bmp/tif/webp/hevc/wxgf) lack a mandatory trailer signature, so they skip the check to avoid false rejection. Also fixes a latent bug in `test_decode_image_v1_no_aes_key_uses_fixed_key`: the test built the synthetic .dat with `TEST_XOR_KEY=0x37` but constructed `ImageResolver` without `xor_key=`, defaulting to `0x88`. The XOR segment was always scrambled — the test passed because the AES segment alone was enough for `detect_image_format` to return 'png' from the header, and no trailer validation existed to catch the corruption. The new trailer check surfaces this, so the test now passes `xor_key=TEST_XOR_KEY` explicitly. Tests: 5 new cases (wrong AES key / wrong XOR for jpg / wrong XOR for png / xor_size=0 bypass / wxgf bypass). All 156 existing tests still pass. --- decode_image.py | 16 +++++++ tests/test_decode_image_v2.py | 83 ++++++++++++++++++++++++++++++++++- 2 files changed, 98 insertions(+), 1 deletion(-) diff --git a/decode_image.py b/decode_image.py index 870bca2..5f7be00 100644 --- a/decode_image.py +++ b/decode_image.py @@ -189,6 +189,22 @@ def v2_decrypt_file(dat_path, out_path=None, aes_key=None, xor_key=0x88): # wxgf (HEVC 裸流) 格式 if decrypted[:4] == b'wxgf': fmt = 'hevc' + elif fmt == 'bin': + # detect_image_format 返回 'bin' = magic 不匹配任何已知图片格式, + # 通常说明 AES key 错(解密后产生随机字节)。拒绝写出无意义的 .bin + # 垃圾文件,让 caller 知道解密失败。 + return None, None + elif xor_size >= 2: + # XOR key 错时 AES/raw 段可能产生合法 magic(看似正常 jpg/png 头), + # 但 XOR 段乱码。用尾部 magic 验证 XOR key 正确性: + # - JPG 必须以 FF D9 (EOI marker) 收尾 + # - PNG 末尾 12 字节必须含 IEND chunk + # 其他格式 (gif/bmp/tif/webp/hevc) 缺乏强制 trailer signature, + # 不做校验以避免误杀。xor_size < 2 时无 XOR 段或样本过小,跳过。 + if fmt == 'jpg' and decrypted[-2:] != b'\xff\xd9': + return None, None + if fmt == 'png' and b'IEND' not in decrypted[-12:]: + return None, None if out_path is None: base = os.path.splitext(dat_path)[0] diff --git a/tests/test_decode_image_v2.py b/tests/test_decode_image_v2.py index 2c4112d..95ee556 100644 --- a/tests/test_decode_image_v2.py +++ b/tests/test_decode_image_v2.py @@ -195,6 +195,82 @@ class TestV2DecryptSynthetic(unittest.TestCase): self.assertEqual(fmt, 'hevc') self.assertTrue(out_path.endswith('.hevc')) + def test_v2_rejects_wrong_aes_key(self): + # AES key 错时 detect_image_format 返回 'bin' (magic 不识别),v2_decrypt_file + # 应拒绝写出 .bin 垃圾文件并返回 (None, None),让 caller 知道解密失败。 + 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)) + + wrong_aes_key = b'wrongkey00000000' + out_path, fmt = v2_decrypt_file( + dat_path, aes_key=wrong_aes_key, xor_key=TEST_XOR_KEY, + ) + self.assertIsNone(out_path) + self.assertIsNone(fmt) + + def test_v2_rejects_wrong_xor_key_jpg_trailer(self): + # JPG 必须以 FF D9 (EOI) 收尾。XOR key 错时尾部 16 字节乱码, + # FF D9 被破坏,触发尾部 magic 校验失败。 + jpg_payload = b'\xff\xd8\xff' + b'\x00' * 83 + b'\xff\xd9' # 88 bytes, FF D9 在末尾 + 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(jpg_payload, aes_size=32, xor_size=16)) + + # 翻转所有 XOR 字节: TEST_XOR_KEY ^ 0xff 保证每字节都错位 + out_path, fmt = v2_decrypt_file( + dat_path, aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY ^ 0xff, + ) + self.assertIsNone(out_path) + self.assertIsNone(fmt) + + def test_v2_rejects_wrong_xor_key_png_iend(self): + # PNG 末尾 12 字节必须含 IEND chunk。XOR key 错时 IEND 被破坏。 + 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 ^ 0xff, + ) + self.assertIsNone(out_path) + self.assertIsNone(fmt) + + def test_v2_skip_xor_validation_when_xor_size_zero(self): + # xor_size < 2 时没有 XOR 段(或样本不足以验证),不应触发尾部 magic 校验。 + # 构造 xor_size=0 的 PNG (整张图都在 AES + raw 段),xor_key 实际不参与解密。 + 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=0)) + + # xor_key 传 0 也应成功 (XOR 段长度 0) + out_path, fmt = v2_decrypt_file( + dat_path, aes_key=TEST_AES_KEY, xor_key=0x00, + ) + self.assertIsNotNone(out_path) + self.assertEqual(fmt, 'png') + + def test_v2_wxgf_skips_trailer_validation(self): + # wxgf (HEVC 裸流) 没有强制 trailer signature,XOR key 错时也不应被尾部校验误杀 + # (wxgf 路径在 elif 链前面命中,直接 fmt='hevc',不进入 XOR 校验分支)。 + # 这里验证:即便 XOR key 错导致末尾字节乱码,只要 wxgf magic 在头部正确, + # 仍按 hevc 输出 — 因为我们只校验 jpg/png,其他格式跳过。 + wxgf_payload = b'wxgf' + b'\x00' * 84 + 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 = v2_decrypt_file( + dat_path, aes_key=TEST_AES_KEY, xor_key=TEST_XOR_KEY ^ 0xff, + ) + self.assertIsNotNone(out_path) + self.assertEqual(fmt, 'hevc') + class TestImageResolverV2(unittest.TestCase): """ImageResolver 端到端:从 local_id 到解密文件,验证 V2 keys 注入路径""" @@ -270,7 +346,12 @@ class TestImageResolverV2(unittest.TestCase): aes_key=v1_fixed_key, magic=V1_MAGIC_FULL, )) - resolver = ImageResolver(self.wechat_base, self.out_dir, self.cache, aes_key=None) + # xor_key 必须跟 _build_v2_dat 加密时用的一致,否则 XOR 段乱码, + # 触发新的尾部 magic 校验失败 (PNG IEND chunk 错位)。 + resolver = ImageResolver( + self.wechat_base, self.out_dir, self.cache, + aes_key=None, 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')