mirror of
https://github.com/LifeArchiveProject/WeChatDataAnalysis.git
synced 2026-02-19 22:30:49 +08:00
5f3c9d82e5
- 新增 sns_media 模块：CDN URL 归一化、远程下载、图片 wcdb_api 解密、视频 WxIsaac64(WeFlow WASM)/ISAAC64 兜底解密与缓存 - routers/sns 与 sns_export_service 复用该模块，收敛重复实现 - 调整 ISAAC64 兜底实现：明确 keystream 生成与字节序格式，作为 WASM 不可用时的 best-effort - 增加单测覆盖：URL 改写、视频异或解密、缓存命中/升级、解密失败
181 lines
7.4 KiB
Python
181 lines
7.4 KiB
Python
import asyncio
|
|
import hashlib
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT / "src"))
|
|
|
|
|
|
from wechat_decrypt_tool import sns_media # noqa: E402 pylint: disable=wrong-import-position
|
|
|
|
|
|
class TestSnsMedia(unittest.TestCase):
|
|
def test_fix_sns_cdn_url_image_rewrites_150_and_appends_token(self):
|
|
u = "http://mmsns.qpic.cn/sns/abc/150"
|
|
out = sns_media.fix_sns_cdn_url(u, token="tkn", is_video=False)
|
|
self.assertEqual(out, "https://mmsns.qpic.cn/sns/abc/0?token=tkn&idx=1")
|
|
|
|
u2 = "https://mmsns.qpic.cn/sns/abc/150?foo=bar"
|
|
out2 = sns_media.fix_sns_cdn_url(u2, token="tkn", is_video=False)
|
|
self.assertEqual(out2, "https://mmsns.qpic.cn/sns/abc/0?foo=bar&token=tkn&idx=1")
|
|
|
|
def test_fix_sns_cdn_url_video_places_token_first(self):
|
|
u = "https://snsvideodownload.video.qq.com/abc.mp4?foo=1&bar=2"
|
|
out = sns_media.fix_sns_cdn_url(u, token="tkn", is_video=True)
|
|
self.assertEqual(out, "https://snsvideodownload.video.qq.com/abc.mp4?token=tkn&idx=1&foo=1&bar=2")
|
|
|
|
def test_fix_sns_cdn_url_non_tencent_host_passthrough(self):
|
|
u = "http://example.com/a/150?x=1"
|
|
out = sns_media.fix_sns_cdn_url(u, token="tkn", is_video=False)
|
|
self.assertEqual(out, u)
|
|
|
|
def test_maybe_decrypt_sns_video_file_xors_inplace(self):
|
|
# Build a fake MP4 header (ftyp at offset 4) and encrypt it by XORing with a keystream.
|
|
plain = b"\x00\x00\x00\x20ftypisom" + b"\x00" * 48
|
|
ks = bytes(range(len(plain)))
|
|
enc = bytes([plain[i] ^ ks[i] for i in range(len(plain))])
|
|
|
|
with TemporaryDirectory() as td:
|
|
p = Path(td) / "v.mp4"
|
|
p.write_bytes(enc)
|
|
|
|
with mock.patch("wechat_decrypt_tool.sns_media.weflow_wxisaac64_keystream", return_value=ks):
|
|
did = sns_media.maybe_decrypt_sns_video_file(p, key="1")
|
|
self.assertTrue(did)
|
|
self.assertEqual(p.read_bytes(), plain)
|
|
|
|
# Second run should be a no-op because it already looks like a MP4.
|
|
did2 = sns_media.maybe_decrypt_sns_video_file(p, key="1")
|
|
self.assertFalse(did2)
|
|
|
|
def test_try_fetch_and_decrypt_sns_image_remote_cache_hit(self):
|
|
with TemporaryDirectory() as td:
|
|
account_dir = Path(td) / "acc"
|
|
account_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
url = "https://mmsns.qpic.cn/sns/test/0?token=tkn&idx=1"
|
|
key = "123"
|
|
fixed = sns_media.fix_sns_cdn_url(url, token="tkn", is_video=False)
|
|
digest = hashlib.md5(f"{fixed}|{key}".encode("utf-8", errors="ignore")).hexdigest()
|
|
|
|
cache_dir = account_dir / "sns_remote_cache" / digest[:2]
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_path = cache_dir / f"{digest}.jpg"
|
|
|
|
payload = b"\xff\xd8\xff\x00fakejpeg"
|
|
cache_path.write_bytes(payload)
|
|
|
|
res = asyncio.run(
|
|
sns_media.try_fetch_and_decrypt_sns_image_remote(
|
|
account_dir=account_dir,
|
|
url=url,
|
|
key=key,
|
|
token="tkn",
|
|
use_cache=True,
|
|
)
|
|
)
|
|
self.assertIsNotNone(res)
|
|
assert res is not None
|
|
self.assertEqual(res.source, "remote-cache")
|
|
self.assertEqual(res.media_type, "image/jpeg")
|
|
self.assertEqual(res.payload, payload)
|
|
self.assertTrue(res.cache_path and res.cache_path.exists())
|
|
|
|
def test_try_fetch_and_decrypt_sns_image_remote_cache_upgrades_bin_extension(self):
|
|
with TemporaryDirectory() as td:
|
|
account_dir = Path(td) / "acc"
|
|
account_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
url = "https://mmsns.qpic.cn/sns/test/0?token=tkn&idx=1"
|
|
key = "123"
|
|
fixed = sns_media.fix_sns_cdn_url(url, token="tkn", is_video=False)
|
|
digest = hashlib.md5(f"{fixed}|{key}".encode("utf-8", errors="ignore")).hexdigest()
|
|
|
|
cache_dir = account_dir / "sns_remote_cache" / digest[:2]
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
bin_path = cache_dir / f"{digest}.bin"
|
|
png_payload = b"\x89PNG\r\n\x1a\n" + b"fakepng"
|
|
bin_path.write_bytes(png_payload)
|
|
|
|
res = asyncio.run(
|
|
sns_media.try_fetch_and_decrypt_sns_image_remote(
|
|
account_dir=account_dir,
|
|
url=url,
|
|
key=key,
|
|
token="tkn",
|
|
use_cache=True,
|
|
)
|
|
)
|
|
self.assertIsNotNone(res)
|
|
assert res is not None
|
|
self.assertEqual(res.source, "remote-cache")
|
|
self.assertEqual(res.media_type, "image/png")
|
|
self.assertTrue(res.cache_path and res.cache_path.suffix.lower() == ".png")
|
|
self.assertTrue(res.cache_path and res.cache_path.exists())
|
|
self.assertFalse(bin_path.exists())
|
|
|
|
def test_try_fetch_and_decrypt_sns_image_remote_decrypts_when_needed(self):
|
|
raw = b"\x01\x02\x03\x04not_an_image"
|
|
decoded = b"\x89PNG\r\n\x1a\n" + b"decoded"
|
|
|
|
async def fake_download(_url: str):
|
|
return raw, "image/jpeg", "1"
|
|
|
|
with TemporaryDirectory() as td:
|
|
account_dir = Path(td) / "acc"
|
|
account_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
with mock.patch("wechat_decrypt_tool.sns_media._download_sns_remote_bytes", side_effect=fake_download):
|
|
with mock.patch("wechat_decrypt_tool.sns_media._wcdb_decrypt_sns_image", return_value=decoded):
|
|
res = asyncio.run(
|
|
sns_media.try_fetch_and_decrypt_sns_image_remote(
|
|
account_dir=account_dir,
|
|
url="https://mmsns.qpic.cn/sns/test/0",
|
|
key="123",
|
|
token="tkn",
|
|
use_cache=False,
|
|
)
|
|
)
|
|
|
|
self.assertIsNotNone(res)
|
|
assert res is not None
|
|
self.assertEqual(res.media_type, "image/png")
|
|
self.assertEqual(res.source, "remote-decrypt")
|
|
self.assertEqual(res.x_enc, "1")
|
|
self.assertEqual(res.payload, decoded)
|
|
|
|
def test_try_fetch_and_decrypt_sns_image_remote_decrypt_failure_returns_none(self):
|
|
raw = b"\x01\x02\x03\x04not_an_image"
|
|
decoded_bad = b"\x00\x00\x00\x00still_bad"
|
|
|
|
async def fake_download(_url: str):
|
|
return raw, "image/jpeg", "1"
|
|
|
|
with TemporaryDirectory() as td:
|
|
account_dir = Path(td) / "acc"
|
|
account_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
with mock.patch("wechat_decrypt_tool.sns_media._download_sns_remote_bytes", side_effect=fake_download):
|
|
with mock.patch("wechat_decrypt_tool.sns_media._wcdb_decrypt_sns_image", return_value=decoded_bad):
|
|
res = asyncio.run(
|
|
sns_media.try_fetch_and_decrypt_sns_image_remote(
|
|
account_dir=account_dir,
|
|
url="https://mmsns.qpic.cn/sns/test/0",
|
|
key="123",
|
|
token="tkn",
|
|
use_cache=False,
|
|
)
|
|
)
|
|
|
|
self.assertIsNone(res)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|