nonebot-plugin-shitbot 0.1.0__py3-none-any.whl

nonebot_plugin_shitbot/commands/convert_cmd.py

@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+import asyncio
+import shutil
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import httpx
+from nonebot.adapters.onebot.v11 import Bot, Message, MessageEvent, MessageSegment
+from nonebot.log import logger
+
+from ..aux import rm_cache, stuff_download
+from ..command import BotCommand
+from ..config import config
+from ..msgutils import dump_message, get_multimedias_url
+from ..parser import BotArgParser
+from ..tasks import EndOfQueue, prod_cons
+
+if TYPE_CHECKING:
+    from ..session import BotSession
+
+bash = shutil.which("bash") or "/bin/bash"
+
+
+class BotCommandConvert(BotCommand):
+    _name = "convert"
+
+    def __init__(self, bot: Bot, session: BotSession, *, _pid: int, _internal=None):
+        super().__init__(bot, session, _pid=_pid, _internal=_internal)
+        self._mode = "video"
+        self._runlock = asyncio.Lock()
+        self._if_accept_pic = False
+        self._prod_lock = asyncio.Lock()
+        self._convert_lock = asyncio.Lock()
+        self._urls: asyncio.Queue[str | EndOfQueue] = asyncio.Queue()
+        self._downloads: asyncio.Queue[str | EndOfQueue] = asyncio.Queue()
+        self._pngs: asyncio.Queue[str | EndOfQueue] = asyncio.Queue()
+        self._outputs: asyncio.Queue[str | EndOfQueue] = asyncio.Queue()
+
+    @property
+    def outputs(self):
+        return self._outputs
+
+    def _init_parser(self):
+        parser = BotArgParser()
+        start = parser.add_subparser("start")
+        stop = parser.add_subparser("stop")
+        start.set_rule(max=0)
+        start.add_opt("-m", required=True, choice=["video", "frame"], default=["video"])
+        stop.set_rule(max=0)
+        parser.set_rule(max=0, need_subcmd=True)
+        return parser
+
+    async def _guard_state(self, subcmd=None):
+        if subcmd is None:
+            return False
+        if self._argv is not None and subcmd == "start":
+            if (
+                self._session is None
+                or (command := self._session.commands.get(self._pid)) is None
+            ):
+                return False
+            tip = "错误：会话被占用\n"
+            tip += f"命令 {command.name} 正在运行，进行下一步前请先终止它或等待其完成。"
+            await self.send_msg(tip)
+            return False
+
+        if self._argv is None and subcmd == "stop":
+            tip = "错误：会话未开始\n"
+            tip += "你还没有开始收集图片，请先使用 /convert start 。"
+            await self.send_msg(tip)
+            self.unlock()
+            return False
+
+        return True
+
+    @staticmethod
+    async def _url_to_download(
+        download_url: str, client: httpx.AsyncClient, download_dir: Path
+    ):
+        try:
+            filename = f"{uuid.uuid4().hex}"
+            save_path = download_dir / filename
+            await stuff_download(client, download_url, save_path)
+            logger.info(f"下载图片成功: {save_path}")
+            return str(save_path)
+        except Exception as e:
+            logger.error(f"下载图片失败 {download_url}: {e}")
+            return ""
+
+    @staticmethod
+    async def _download_to_png(download_path: str, png_dir: Path):
+        if download_path == "":
+            return ""
+        png_path = png_dir / f"{uuid.uuid4().hex}.png"
+        logger.info(
+            f"执行图片转png脚本: {config.script_p2png_path} {download_path} {png_path}"
+        )
+        proc = await asyncio.create_subprocess_exec(
+            bash,
+            str(config.script_p2png_path),
+            str(download_path),
+            str(png_path),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await proc.communicate()
+        if stdout:
+            logger.info(f"脚本 stdout:\n{stdout.decode()}")
+        if stderr:
+            logger.warning(f"脚本 stderr:\n{stderr.decode()}")
+        if proc.returncode != 0:
+            logger.error(f"转换失败: {stderr.decode()[:]}")
+            return ""
+        logger.info(f"转换图片成功: {png_path}")
+        return str(png_path)
+
+    @staticmethod
+    async def _png_to_video(png_path: str, video_dir: Path):
+        if png_path == "":
+            return ""
+        video_path = video_dir / f"{uuid.uuid4().hex}.mp4"
+        logger.info(
+            f"执行图片转视频脚本: {config.script_png2v_path} {png_path} {video_path}"
+        )
+        proc = await asyncio.create_subprocess_exec(
+            bash,
+            str(config.script_png2v_path),
+            str(png_path),
+            str(video_path),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await proc.communicate()
+        if stdout:
+            logger.info(f"脚本 stdout:\n{stdout.decode()}")
+        if stderr:
+            logger.warning(f"脚本 stderr:\n{stderr.decode()}")
+        if proc.returncode != 0:
+            logger.error(f"转换失败: {stderr.decode()[:]}")
+            return ""
+        logger.info(f"转换视频成功: {video_path}")
+        return str(video_path)
+
+    @staticmethod
+    async def _png_to_frame(png_path: str, frame_dir: Path):
+        if png_path == "":
+            return ""
+        frame_path = frame_dir / f"{uuid.uuid4().hex}.png"
+        logger.info(
+            f"执行图片加框脚本: {config.script_png2fr_path} {png_path} {frame_path}"
+        )
+        proc = await asyncio.create_subprocess_exec(
+            bash,
+            str(config.script_png2fr_path),
+            str(png_path),
+            str(frame_path),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await proc.communicate()
+        if stdout:
+            logger.info(f"脚本 stdout:\n{stdout.decode()}")
+        if stderr:
+            logger.warning(f"脚本 stderr:\n{stderr.decode()}")
+        if proc.returncode != 0:
+            logger.error(f"转换失败: {stderr.decode()[:]}")
+            return ""
+        logger.info(f"转换方形图片成功: {frame_path}")
+        return str(frame_path)
+
+    async def _send_outputs(self):
+        if not self.session:
+            return
+        outputs = []
+        while True:
+            output = await self._outputs.get()
+            if isinstance(output, EndOfQueue):
+                break
+            if output == "":
+                continue
+            outputs.append(output)
+
+        logger.info(f"用户 {self.session.user_id} 结束收集，共收到 {len(outputs)} 张")
+
+        if len(outputs) == 0:
+            await self.send_msg("没有输出被生成。")
+            return
+
+        await self.send_msg(f"共生成 {len(outputs)} 个输出，逐个发送…")
+
+        for i, output in enumerate(outputs):
+            output_file = Path(output)
+            if not output_file.exists():
+                logger.error(f"输出文件不存在: {output_file}")
+                continue
+            container_path = config.client_base / output_file.relative_to(
+                config.bot_base
+            )
+            try:
+                if self._mode == "video":
+                    await self.send_msg(
+                        Message(MessageSegment.video(f"file://{container_path}"))
+                    )
+                elif self._mode == "frame":
+                    await self.send_msg(
+                        Message(MessageSegment.image(f"file://{container_path}"))
+                    )
+                logger.info(f"发送第 {i + 1} 个输出成功: {output_file.name}")
+            except Exception as e:
+                logger.error(f"发送 {output_file.name} 失败: {e}")
+                await self.send_msg(
+                    f"发送 {output_file.name} 失败: {e}\n尝试以文件形式发送..."
+                )
+                try:
+                    msg = Message(
+                        MessageSegment("file", {"file": f"file://{container_path}"})
+                    )
+                    await self.send_msg(msg)
+                    logger.info(
+                        f"以文件形式发送第 {i + 1} 个输出成功: {output_file.name}"
+                    )
+                except Exception as e:
+                    logger.error(f"发送 {output_file.name} 失败: {e}")
+                    await self.send_msg(
+                        f"发送第 {i + 1} 个输出 {output_file.name} 失败: {e}"
+                    )
+
+                await asyncio.sleep(0.25)
+
+        await self.send_msg("输出发送完毕。")
+
+    async def _convert_start(self):
+        if not self.session:
+            return
+        downloads_dir = (
+            config.cache
+            / self.session.group_id
+            / self.session.user_id
+            / str(self._pid)
+            / "downloads"
+        )
+        pngs_dir = (
+            config.cache
+            / self.session.group_id
+            / self.session.user_id
+            / str(self._pid)
+            / "pngs"
+        )
+        outputs_dir = (
+            config.cache
+            / self.session.group_id
+            / self.session.user_id
+            / str(self._pid)
+            / "outputs"
+        )
+
+        downloads_dir.mkdir(parents=True, exist_ok=True)
+        pngs_dir.mkdir(parents=True, exist_ok=True)
+        outputs_dir.mkdir(parents=True, exist_ok=True)
+
+        async def _urls_to_downloads():
+            try:
+                async with httpx.AsyncClient() as client:
+                    await prod_cons(
+                        self._urls,
+                        self._downloads,
+                        self._url_to_download,
+                        client,
+                        downloads_dir,
+                    )
+            except Exception as e:
+                logger.exception(f"urls_to_downloads 管道异常退出: {e}")
+
+        async def _downloads_to_pngs():
+            try:
+                await prod_cons(
+                    self._downloads, self._pngs, self._download_to_png, pngs_dir
+                )
+            except Exception as e:
+                logger.exception(f"downloads_to_pngs 管道异常退出: {e}")
+
+        async def _pngs_to_outputs():
+            try:
+                async with self._convert_lock:
+                    if self._mode == "video":
+                        await prod_cons(
+                            self._pngs, self._outputs, self._png_to_video, outputs_dir
+                        )
+                    elif self._mode == "frame":
+                        await prod_cons(
+                            self._pngs, self._outputs, self._png_to_frame, outputs_dir
+                        )
+            except Exception as e:
+                logger.exception(f"pngs_to_outputs 管道异常退出: {e}")
+
+        asyncio.create_task(_urls_to_downloads())
+        asyncio.create_task(_downloads_to_pngs())
+        asyncio.create_task(_pngs_to_outputs())
+
+        logger.info(f"用户 {self.session.user_id} 开始了图片收集")
+        await self.send_msg(
+            "图片收集已开始， Bot 会收集本条信息后你发送的所有图片，直到你发送 /convert stop 完成收集。"
+        )
+
+        self._if_accept_pic = True
+        return
+
+    async def _convert_stop(self):
+        if not self.session:
+            return
+        self._if_accept_pic = False
+
+        await self.send_msg("转换图片中...")
+        async with self._prod_lock:
+            await self._urls.put(EndOfQueue())
+
+        async with self._convert_lock:
+            pass
+        await self.send_msg("转换完毕。")
+
+        await self._send_outputs()
+        await rm_cache(self.session.group_id, self.session.user_id, str(self._pid))
+        self.unlock()
+        return
+
+    async def roger(self, event: MessageEvent):
+        async with self._prod_lock:
+            if not self._if_accept_pic:
+                return
+            msg = await dump_message(self.bot, event.get_message())
+            url_list = get_multimedias_url(msg, basetypes=["image"])
+            for url in url_list:
+                await self._urls.put(url)
+
+    async def run(self, args: Message):
+        async with self._runlock:
+            if not self.session:
+                return
+
+            new_argv = args.extract_plain_text().strip().split()
+            if not await self._legal_case(new_argv):
+                if self._argv is None:
+                    self.unlock()
+                return
+
+            self._parser.parse_argv(new_argv)
+            subcmd = self._parser.subcmd
+
+            if not await self._guard_state(subcmd):
+                return
+
+            self._argv = new_argv
+
+            if self._pid >= 0 and not self._check_perm("convert"):
+                await self.send_msg("权限不足")
+                self.unlock()
+                return
+
+            if subcmd == "start":
+                new_mode = self._parser.subparsers[subcmd].opts_value.get(
+                    "-m", ["video"]
+                )[0]
+                self._mode = new_mode
+                await self._convert_start()
+                return
+
+            if subcmd == "stop":
+                await self._convert_stop()
+                return

nonebot_plugin_shitbot/commands/help_cmd.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from nonebot.adapters.onebot.v11 import Bot, Message
+
+from ..command import BotCommand
+from ..parser import BotArgParser
+from .md2pic_cmd import BotCommandMd2pic
+
+if TYPE_CHECKING:
+    from ..session import BotSession
+
+
+class BotCommandHelp(BotCommand):
+    _name = "help"
+
+    def __init__(self, bot: Bot, session: BotSession, *, _pid: int, _internal=None):
+        super().__init__(bot, session, _pid=_pid, _internal=_internal)
+
+    def _init_parser(self):
+        parser = BotArgParser()
+        parser.add_subparser("help")
+        parser.add_subparser("session")
+        parser.add_subparser("perm")
+        parser.add_subparser("convert")
+        parser.add_subparser("randpic")
+        parser.add_subparser("shitpost")
+        parser.add_subparser("advrandpic")
+        parser.add_subparser("md2pic")
+        parser.add_subparser("pixiv")
+        return parser
+
+    async def run(self, args: Message):
+        if not self.session:
+            return
+        new_argv = args.extract_plain_text().strip().split()
+        # I know it's impossible to be illegal. Just for formalism :)
+        if not await self._legal_case(new_argv):
+            if self._argv is None:
+                self.unlock()
+            return
+
+        self._argv = new_argv
+        self._parser.parse_argv(self._argv)
+        subcmd = self._parser.subcmd
+
+        if not self._check_perm("help"):
+            await self.send_msg("权限不足")
+            self.unlock()
+            return
+
+        help_dir = Path(__file__).resolve().parent.parent / "docs" / "help"
+        tip = "错误: 帮助文档目录不存在"
+        if help_dir.exists():
+            tip = "错误: 帮助文档 index.md 不存在"
+        help_path = help_dir / "index.md"
+        if help_path.exists():
+            tip = help_path.read_text(encoding="utf-8")
+        if subcmd in (
+            "help",
+            "session",
+            "perm",
+            "convert",
+            "randpic",
+            "shitpost",
+            "advrandpic",
+            "md2pic",
+            "pixiv",
+        ):
+            tip = f"错误: 帮助文档 {subcmd}.md 不存在"
+            help_path = help_dir / f"{subcmd}.md"
+            if help_path.exists():
+                tip = help_path.read_text(encoding="utf-8")
+
+        tip = re.sub(
+            r"!\[([^\]]*)\]\((?!https?://|/)([^)]+)\)",
+            lambda m: f"![{m.group(1)}]({(help_dir / m.group(2)).resolve()})",
+            tip,
+        )
+
+        _pid = -1
+        pids = self.session.commands.keys()
+        while _pid in pids:
+            _pid -= 1
+        md2pic = BotCommandMd2pic.make(self.bot, self.session, _pid=_pid)
+        tip = "-c\n" + tip
+        msg = Message(tip)
+        await md2pic.run(msg)  # type: ignore
+        self.unlock()

nonebot_plugin_shitbot/commands/md2pic_cmd.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import io
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from nonebot import require
+from nonebot.adapters.onebot.v11 import Bot, Message, MessageSegment
+from PIL import Image
+
+require("nonebot_plugin_htmlrender")
+from nonebot_plugin_htmlrender import md_to_pic
+
+from ..command import BotCommand
+from ..parser import BotArgParser
+
+if TYPE_CHECKING:
+    from ..session import BotSession
+
+
+class BotCommandMd2pic(BotCommand):
+    _name = "md2pic"
+
+    def __init__(self, bot: Bot, session: BotSession, *, _pid: int, _internal=None):
+        super().__init__(bot, session, _pid=_pid, _internal=_internal)
+        self._md = ""
+        self._scale = 2
+        self._css = ""
+        self._padding = 30
+        self._min_w = 20
+        self._max_w = 2000
+        self._min_h = 20
+        self._max_h = None
+
+    def _init_parser(self):
+        parser = BotArgParser()
+        parser.set_rule(min=1, max=1)
+        parser.add_opt("-c", necessary=True)
+        parser.add_opt("-s", required=True, type=float, default=[2])
+        parser.add_opt(
+            "-t",
+            required=True,
+            choice=["github-markdown-dark-dimmed"],
+            default=["github-markdown-dark-dimmed"],
+        )
+        parser.add_opt("--padding", required=True, type=int, default=[30])
+        parser.add_opt("--min_w", required=True, type=int, default=[20])
+        parser.add_opt("--max_w", required=True, type=int, default=[2000])
+        parser.add_opt("--min_h", required=True, type=int, default=[20])
+        parser.add_opt("--max_h", required=True, type=int, default=[])
+
+        return parser
+
+    async def _render(self) -> bytes:
+
+        min_w = int(self._scale * self._min_w)
+        max_w = self._max_w
+        min_h = int(self._scale * self._min_h)
+        max_h = (
+            int(self._scale * self._max_h) if self._max_h is not None else self._max_h
+        )
+        padding = int(self._scale * self._padding)
+
+        raw = await md_to_pic(
+            self._md, width=max_w, css_path=self._css, device_scale_factor=self._scale
+        )
+        img = Image.open(io.BytesIO(raw))
+
+        # Get the background color (the upper-left color)
+        bgpx: Any = img.getpixel((0, 0))
+        bgr, bgg, bgb = int(bgpx[0]), int(bgpx[1]), int(bgpx[2])
+
+        def is_bg(r: int, g: int, b: int, *, delta: int = 6) -> bool:
+            return (
+                abs(r - bgr) < delta and abs(g - bgg) < delta and abs(b - bgb) < delta
+            )
+
+        def is_column_has_content(x: int):
+            for y in range(0, h, 4):
+                pixel: Any = img.getpixel((x, y))
+                r, g, b = int(pixel[0]), int(pixel[1]), int(pixel[2])
+                if not is_bg(r, g, b):
+                    return True
+            return False
+
+        def is_row_has_content(y: int):
+            for x in range(0, w, 4):
+                pixel: Any = img.getpixel((x, y))
+                r, g, b = int(pixel[0]), int(pixel[1]), int(pixel[2])
+                if not is_bg(r, g, b):
+                    return True
+            return False
+
+        w, h = img.size
+        content_left, content_right = 0, w
+        content_top, content_bottom = 0, h
+
+        # Find out the leftmost column that has content
+        for x in range(w):
+            if is_column_has_content(x):
+                content_left = x
+                break
+
+        # Find out the rightmost column that has content
+        for x in range(w - 1, content_left - 1, -1):
+            if is_column_has_content(x):
+                content_right = x + 1
+                break
+
+        # Find out the top row that has content
+        for y in range(h):
+            if is_row_has_content(y):
+                content_top = y
+                break
+
+        # Find out the bottom row that has content
+        for y in range(h - 1, content_top - 1, -1):
+            if is_row_has_content(y):
+                content_bottom = y + 1
+                break
+
+        crop_left = max(0, content_left - padding)
+        crop_right = max(crop_left + min_w + 1, content_right + padding)
+        crop_top = max(0, content_top - padding)
+        crop_bottom = max(crop_top + min_h + 1, content_bottom + padding)
+        if max_h is not None:
+            crop_bottom = min(crop_bottom, crop_top + max_h)
+
+        crop_right = min(crop_right, w)
+        crop_bottom = min(crop_bottom, h)
+        img = img.crop((crop_left, crop_top, crop_right, crop_bottom))
+
+        buf = io.BytesIO()
+        img.save(buf, format="PNG")
+        return buf.getvalue()
+
+    async def run(self, args: Message):
+        if not self.session:
+            return
+        new_argss = args.extract_plain_text().split("\n", 1)
+        if len(new_argss) < 2:
+            await self._send_format_error()
+            if self._argv is None:
+                self.unlock()
+            return
+        new_argv = new_argss[0].strip().split()
+        new_argv.append(new_argss[1])
+        if not await self._legal_case(new_argv):
+            if self._argv is None:
+                self.unlock()
+            return
+
+        if not await self._guard_state():
+            return
+
+        self._argv = new_argv
+        self._parser.parse_argv(self._argv)
+
+        if self._pid >= 0 and not self._check_perm("md2pic"):
+            await self.send_msg("权限不足")
+            self.unlock()
+            return
+
+        self._scale = self._parser.opts_value["-s"][0]
+        theme = self._parser.opts_value["-t"][0]
+        self._css = str(Path(__file__).resolve().parent.parent / "css" / f"{theme}.css")
+        self._padding = self._parser.opts_value["--padding"][0]
+        self._min_w = self._parser.opts_value["--min_w"][0]
+        self._max_w = self._parser.opts_value["--max_w"][0]
+        self._min_h = self._parser.opts_value["--min_h"][0]
+        self._max_h = (
+            self._parser.opts_value["--max_h"][0]
+            if len(self._parser.opts_value["--max_h"]) > 0
+            else None
+        )
+        self._md = self._parser.value[0]
+
+        self._padding = max(0, self._padding)
+        self._padding = min(50, self._padding)
+        self._max_w = min(10000000, self._max_w)
+        self._max_w = max(0, self._max_w)
+        self._min_w = min(self._max_w, self._min_w)
+        self._min_w = max(0, self._min_w)
+        if self._max_h is not None:
+            self._max_h = min(10000000, self._max_h)
+            self._max_h = max(0, self._max_h)
+            self._min_h = min(self._max_h, self._min_h)
+        self._min_h = max(0, self._min_h)
+
+        img = await self._render()
+
+        msg = Message(MessageSegment.image(img))
+        await self.send_msg(msg)
+
+        self.unlock()