HeXi/hexi/plugins/nonebot_plugin_cherry_tomato/handlers/obfuscate.py

"""混淆/解混淆 —— 触发器（Trigger）层。

回复一张图片后发「混淆」或「解混淆」触发（图片与命令同一条消息直接发也行，
不需要 @bot）；下载图片 → 变换 → 结果图发回。算法层见 utils/obfuscation.py。
"""

from __future__ import annotations

import re

from nonebot import on_command
from nonebot.adapters.onebot.v11 import Bot, MessageEvent
from nonebot.log import logger
from nonebot_plugin_alconna import UniMessage

from hexi.core.cooldown import cooldown
from hexi.core.message_utils import get_reply_message

from .. import config
from ..services.pipeline import download_image, run_pipeline

# 两个命令共用同一冷却器（name 相同）：同一用户 10 秒内不能连发混淆/解混淆
obf_cmd = cooldown(10, hint="处理太频繁了，{secs} 秒后再试", name="cherry_tomato")(
    on_command("混淆", aliases={"小番茄混淆"}, priority=10, block=True)
)
deobf_cmd = cooldown(10, hint="处理太频繁了，{secs} 秒后再试", name="cherry_tomato")(
    on_command("解混淆", aliases={"小番茄解混淆"}, priority=10, block=True)
)

_MD5_RE = re.compile(r"[0-9A-Fa-f]{32}")


@obf_cmd.handle()
async def _obfuscate(event: MessageEvent, bot: Bot) -> None:
    await _handle(event, bot, decrypt=False)


@deobf_cmd.handle()
async def _deobfuscate(event: MessageEvent, bot: Bot) -> None:
    await _handle(event, bot, decrypt=True)


async def _extract_image(
    event: MessageEvent, bot: Bot
) -> tuple[str | None, str | None]:
    """取被回复（或当前）消息里第一张图的 (url, file)。

    优先段里的 url；没有 url 时用 file 调 OneBot get_image 现换一个。
    file 仅用于 offpic 链接失效时的 md5 兜底重试。
    """
    for seg in get_reply_message(event):
        if seg.type != "image":
            continue
        data = seg.data or {}
        url = data.get("url")
        if url:
            return str(url), str(data.get("file") or "")
        file = data.get("file")
        if not file:
            continue
        try:
            info = await bot.get_image(file=str(file))
        except Exception as e:  # noqa: BLE001
            logger.warning(f"cherry_tomato: get_image 取 url 失败: {e!r}")
            continue
        url = info.get("url") if isinstance(info, dict) else None
        if url:
            return str(url), str(file)
    return None, None


def _offpic_md5(url: str, file: str | None) -> str | None:
    """从 file 字段或 offpic 链接里找 32 位十六进制（md5）串；找不到返回 None。"""
    for candidate in (file, url):
        match = _MD5_RE.search(candidate or "")
        if match:
            return match.group(0).upper()
    return None


async def _download_with_fallback(url: str, file: str | None) -> bytes | None:
    """下载图片；offpic 老链接失败时用 md5 换 gchat 链接重试一次（picfinder 先例）。"""
    data = await download_image(url)
    if data is not None:
        return data
    if "c2cpicdw.qpic.cn/offpic_new/" in url:
        md5 = _offpic_md5(url, file)
        if md5:
            retry_url = f"http://gchat.qpic.cn/gchatpic_new/0/0-0-{md5}/0?term=2"
            logger.info("cherry_tomato: offpic 原链下载失败，改用 gchat 链接重试")
            data = await download_image(retry_url)
    return data


async def _handle(event: MessageEvent, bot: Bot, *, decrypt: bool) -> None:
    action = "解混淆" if decrypt else "混淆"

    url, file = await _extract_image(event, bot)
    if not url:
        await UniMessage.text(
            "请回复一张图片再发这个命令，或把图片和命令一起发~"
        ).finish()

    data = await _download_with_fallback(url, file)
    if data is None:
        logger.warning(f"cherry_tomato: 图片下载失败 url={url}")
        await UniMessage.text("图片下载失败，可能链接已过期，重发一次试试").finish()

    logger.info(f"cherry_tomato: {action}开始 user={event.get_user_id()} url={url}")
    max_pixels = config.decrypt_max_pixels() if decrypt else config.max_pixels()
    out, err, notes = await run_pipeline(
        data, decrypt=decrypt, max_pixels=max_pixels, quality=config.jpeg_quality()
    )
    if out is None:
        logger.warning(f"cherry_tomato: {action}失败: {err}")
        await UniMessage.text(err or f"{action}失败了，换一张图试试").finish()

    image = UniMessage.image(raw=out)
    if notes:
        await (UniMessage.text("；".join(notes) + "\n") + image).send()
    else:
        await image.send()