sansenhoshiandClaude Code
fb15b976aa
Gilbert 曲线像素重排，可与网页版互解；`混淆` / `解混淆` 命令 （回复图片或图文同条消息发送），附 tests/test_cherry_tomato.py Co-Authored-By: Claude Code <noreply@anthropic.com>
293 lines
9.4 KiB
Python
293 lines
9.4 KiB
Python
"""小番茄图片混淆 —— 纯算法测试。
|
||
|
||
utils/obfuscation.py 无 nonebot 依赖，用 importlib 按文件路径裸加载
|
||
（先例 test_rate_limit.py），避免拉起 NoneBot 运行时。
|
||
|
||
核心断言：
|
||
- `_gilbert_positions` 与参考实现（逐点 append 版 gilbert2d，原样抄录）逐元素
|
||
一致 —— 这是「与网页版互解」的关键（曲线必须同一个）；
|
||
- 位置序列恒为 0..total-1 的排列；
|
||
- encrypt / decrypt 互为逆操作（数组级 + PIL 级），偏移公式与参考一致。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import importlib.util
|
||
import io
|
||
import math
|
||
import random
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import pytest
|
||
from PIL import Image
|
||
|
||
_MODULE_PATH = (
|
||
Path(__file__).resolve().parents[1]
|
||
/ "hexi"
|
||
/ "plugins"
|
||
/ "nonebot_plugin_cherry_tomato"
|
||
/ "utils"
|
||
/ "obfuscation.py"
|
||
)
|
||
|
||
|
||
@pytest.fixture(scope="module")
|
||
def obf():
|
||
"""以独立模块名加载 obfuscation.py，避免触发插件包 __init__ 的依赖链。"""
|
||
spec = importlib.util.spec_from_file_location(
|
||
"hexi_cherry_tomato_obfuscation_test", _MODULE_PATH
|
||
)
|
||
assert spec is not None and spec.loader is not None
|
||
module = importlib.util.module_from_spec(spec)
|
||
sys.modules[spec.name] = module
|
||
spec.loader.exec_module(module)
|
||
yield module
|
||
sys.modules.pop(spec.name, None)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────
|
||
# 参考实现（Cherry-tomato-image-obfuscation 逐点 append 版），
|
||
# 原样抄录，仅作对照，不执行于生产路径
|
||
# ─────────────────────────────────────────────────────────────
|
||
|
||
|
||
def _naive_gilbert2d(width: int, height: int) -> list[tuple[int, int]]:
|
||
coordinates: list[tuple[int, int]] = []
|
||
if width >= height:
|
||
_naive_generate2d(0, 0, width, 0, 0, height, coordinates)
|
||
else:
|
||
_naive_generate2d(0, 0, 0, height, width, 0, coordinates)
|
||
return coordinates
|
||
|
||
|
||
def _naive_sign(x):
|
||
if x > 0:
|
||
return 1
|
||
elif x < 0:
|
||
return -1
|
||
return 0
|
||
|
||
|
||
def _naive_generate2d(x, y, ax, ay, bx, by, coordinates):
|
||
w = abs(ax + ay)
|
||
h = abs(bx + by)
|
||
|
||
dax = _naive_sign(ax)
|
||
day = _naive_sign(ay)
|
||
dbx = _naive_sign(bx)
|
||
dby = _naive_sign(by)
|
||
|
||
if h == 1:
|
||
for _ in range(w):
|
||
coordinates.append((x, y))
|
||
x += dax
|
||
y += day
|
||
return
|
||
|
||
if w == 1:
|
||
for _ in range(h):
|
||
coordinates.append((x, y))
|
||
x += dbx
|
||
y += dby
|
||
return
|
||
|
||
ax2 = ax // 2
|
||
ay2 = ay // 2
|
||
bx2 = bx // 2
|
||
by2 = by // 2
|
||
|
||
w2 = abs(ax2 + ay2)
|
||
h2 = abs(bx2 + by2)
|
||
|
||
if 2 * w > 3 * h:
|
||
if (w2 % 2) and (w > 2):
|
||
ax2 += dax
|
||
ay2 += day
|
||
_naive_generate2d(x, y, ax2, ay2, bx, by, coordinates)
|
||
_naive_generate2d(x + ax2, y + ay2, ax - ax2, ay - ay2, bx, by, coordinates)
|
||
else:
|
||
if (h2 % 2) and (h > 2):
|
||
bx2 += dbx
|
||
by2 += dby
|
||
_naive_generate2d(x, y, bx2, by2, ax2, ay2, coordinates)
|
||
_naive_generate2d(x + bx2, y + by2, ax, ay, bx - bx2, by - by2, coordinates)
|
||
_naive_generate2d(
|
||
x + (ax - dax) + (bx2 - dbx),
|
||
y + (ay - day) + (by2 - dby),
|
||
-bx2,
|
||
-by2,
|
||
-(ax - ax2),
|
||
-(ay - ay2),
|
||
coordinates,
|
||
)
|
||
|
||
|
||
def _naive_positions(width: int, height: int) -> list[int]:
|
||
return [x + y * width for x, y in _naive_gilbert2d(width, height)]
|
||
|
||
|
||
def _build_sizes() -> list[tuple[int, int]]:
|
||
sizes: list[tuple[int, int]] = [
|
||
(1, 1),
|
||
(1, 2),
|
||
(2, 1),
|
||
(1, 3),
|
||
(3, 1),
|
||
(2, 2),
|
||
(3, 3),
|
||
(7, 5),
|
||
(5, 7),
|
||
(1, 33),
|
||
(33, 1),
|
||
(1, 400),
|
||
(400, 1),
|
||
(128, 96),
|
||
(200, 50),
|
||
(50, 200),
|
||
(333, 37),
|
||
(256, 256),
|
||
]
|
||
primes = (2, 3, 5, 13, 31, 61, 79)
|
||
sizes += [(p, q) for p in primes for q in primes]
|
||
rng = random.Random(20261010)
|
||
sizes += [(rng.randint(1, 80), rng.randint(1, 80)) for _ in range(220)]
|
||
seen: set[tuple[int, int]] = set()
|
||
unique: list[tuple[int, int]] = []
|
||
for size in sizes:
|
||
if size not in seen:
|
||
seen.add(size)
|
||
unique.append(size)
|
||
return unique
|
||
|
||
|
||
_SIZES = _build_sizes()
|
||
|
||
|
||
@pytest.mark.parametrize("size", _SIZES, ids=[f"{w}x{h}" for w, h in _SIZES])
|
||
def test_positions_match_reference(obf, size):
|
||
"""任何尺寸下位置序列都与参考逐点实现完全一致（互解的前提）。"""
|
||
width, height = size
|
||
pos = obf._gilbert_positions(width, height)
|
||
assert pos.dtype == np.int32
|
||
assert pos.tolist() == _naive_positions(width, height)
|
||
|
||
|
||
@pytest.mark.parametrize("size", _SIZES, ids=[f"{w}x{h}" for w, h in _SIZES])
|
||
def test_positions_is_permutation(obf, size):
|
||
"""位置序列必须是 0..total-1 的排列（不重不漏）。"""
|
||
width, height = size
|
||
pos = obf._gilbert_positions(width, height)
|
||
assert sorted(pos.tolist()) == list(range(width * height))
|
||
|
||
|
||
def test_gilbert_offset_formula(obf):
|
||
"""黄金比例偏移与参考/JS 公式一致。"""
|
||
for total in (1, 2, 3, 6, 10, 1920 * 1080, 8_000_000):
|
||
assert obf.gilbert_offset(total) == round((math.sqrt(5) - 1) / 2 * total)
|
||
assert obf.gilbert_offset(10) == 6 # 0.618...*10 = 6.18 → 6（与 JS 一致）
|
||
|
||
|
||
@pytest.mark.parametrize("size", [(3, 2), (17, 5), (40, 33), (64, 64)])
|
||
def test_encrypt_decrypt_array_roundtrip(obf, size):
|
||
"""数组级 encrypt∘decrypt 恒等；encrypt 不改输入、确实重排了像素。"""
|
||
width, height = size
|
||
rng = np.random.default_rng(20261010)
|
||
src = rng.integers(0, 256, size=(width * height, 3), dtype=np.uint8)
|
||
original = src.copy()
|
||
|
||
enc = obf.encrypt_pixels(src, width, height)
|
||
assert np.array_equal(src, original) # 输入未被改写
|
||
assert not np.array_equal(enc, src) # 偏移非 0，随机图必被重排
|
||
|
||
dec = obf.decrypt_pixels(enc, width, height)
|
||
assert np.array_equal(dec, src)
|
||
|
||
|
||
@pytest.mark.parametrize("size", [(3, 2), (13, 9), (40, 33)])
|
||
def test_encrypt_decrypt_pil_roundtrip(obf, size):
|
||
"""PIL 级 round-trip：obfuscate 输出为 RGB、无 EXIF，解混淆还原原图。"""
|
||
width, height = size
|
||
rng = np.random.default_rng(7)
|
||
arr = rng.integers(0, 256, (height, width, 3), dtype=np.uint8)
|
||
img = Image.fromarray(arr, "RGB")
|
||
|
||
enc, scaled, extra = obf.obfuscate(img, decrypt=False, max_pixels=10**9)
|
||
assert enc.mode == "RGB"
|
||
assert scaled is False
|
||
assert extra is False
|
||
assert not np.array_equal(np.asarray(enc), arr)
|
||
|
||
dec, _, _ = obf.obfuscate(enc, decrypt=True, max_pixels=10**9)
|
||
assert np.array_equal(np.asarray(dec), arr)
|
||
assert not dec.getexif() # 输出统一不带 EXIF
|
||
|
||
# JPEG 编码可解码、尺寸不变
|
||
data = obf.to_jpeg_bytes(dec, 95)
|
||
with Image.open(io.BytesIO(data)) as reopened:
|
||
assert reopened.size == (width, height)
|
||
|
||
|
||
def test_obfuscate_applies_exif_orientation(obf):
|
||
"""带 EXIF 旋转的照片先规范化方向：输出宽高对调（显示方向不突变）。"""
|
||
img = Image.new("RGB", (13, 9), (10, 20, 30))
|
||
exif = Image.Exif()
|
||
exif[0x0112] = 6 # Orientation=6：显示时需顺时针旋转 90°（宽高对调）
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="JPEG", exif=exif)
|
||
buf.seek(0)
|
||
with Image.open(buf) as photo:
|
||
assert photo.size == (13, 9)
|
||
assert photo.getexif().get(0x0112) == 6
|
||
out, _, _ = obf.obfuscate(photo, decrypt=False, max_pixels=10**9)
|
||
assert out.size == (9, 13)
|
||
|
||
|
||
def test_fit_within(obf):
|
||
"""fit_within：未超限原样返回；超限等比缩小且至少 1 像素。"""
|
||
img = Image.new("RGB", (4000, 3000))
|
||
assert obf.fit_within(img, 12_000_000) is img
|
||
|
||
out = obf.fit_within(img, 2_000_000)
|
||
assert out.width * out.height <= 2_000_000
|
||
assert abs(out.width / out.height - 4000 / 3000) < 0.01 # 等比
|
||
|
||
tiny = obf.fit_within(img, 1)
|
||
assert tiny.width >= 1 and tiny.height >= 1
|
||
|
||
|
||
def test_obfuscate_scale_flag_only_for_encrypt(obf):
|
||
"""加密超限缩放并置 scaled；解密绝不缩放（超大由 pipeline 拒绝兜底）。"""
|
||
img = Image.new("RGB", (60, 40))
|
||
|
||
enc, scaled, _ = obf.obfuscate(img, decrypt=False, max_pixels=1000)
|
||
assert scaled is True
|
||
assert enc.width * enc.height <= 1000
|
||
|
||
dec, scaled2, _ = obf.obfuscate(img, decrypt=True, max_pixels=1000)
|
||
assert scaled2 is False
|
||
assert dec.size == img.size
|
||
|
||
|
||
def test_obfuscate_extra_frames_flag(obf):
|
||
"""多帧图只取第一帧并置 extra_frames；单帧图不置。"""
|
||
frames = [
|
||
Image.new("RGB", (8, 6), (255, 0, 0)),
|
||
Image.new("RGB", (8, 6), (0, 0, 255)),
|
||
]
|
||
buf = io.BytesIO()
|
||
frames[0].save(buf, format="GIF", save_all=True, append_images=frames[1:])
|
||
buf.seek(0)
|
||
with Image.open(buf) as gif:
|
||
out, _, extra = obf.obfuscate(gif, decrypt=False, max_pixels=10**9)
|
||
assert extra is True
|
||
assert out.size == (8, 6)
|
||
# 首帧是纯红：均匀图重排后仍是纯红（同时证明没拿第二帧的蓝色）
|
||
assert out.getpixel((0, 0)) == (255, 0, 0)
|
||
|
||
_, _, extra2 = obf.obfuscate(
|
||
Image.new("RGB", (8, 6)), decrypt=False, max_pixels=10**9
|
||
)
|
||
assert extra2 is False
|