From 27c9394b0da5432b6ab0bc95098d4f44069cd61e Mon Sep 17 00:00:00 2001 From: HibiKier <775757368@qq.com> Date: Sun, 26 May 2024 15:22:55 +0800 Subject: [PATCH] =?UTF-8?q?feat=E2=9C=A8:=20=E8=89=B2=E5=9B=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .vscode/settings.json | 2 + poetry.lock | 81 +++- pyproject.toml | 1 + zhenxun/plugins/send_setu_/__init__.py | 5 + zhenxun/plugins/send_setu_/_model.py | 87 +++++ .../plugins/send_setu_/send_setu/__init__.py | 227 +++++++++++ .../send_setu_/send_setu/_data_source.py | 362 ++++++++++++++++++ .../send_setu_/update_setu/__init__.py | 59 +++ .../send_setu_/update_setu/data_source.py | 187 +++++++++ zhenxun/services/db_context.py | 64 ++-- zhenxun/utils/image_utils.py | 27 ++ zhenxun/utils/withdraw_manage.py | 3 +- 12 files changed, 1071 insertions(+), 34 deletions(-) create mode 100644 zhenxun/plugins/send_setu_/__init__.py create mode 100644 zhenxun/plugins/send_setu_/_model.py create mode 100644 zhenxun/plugins/send_setu_/send_setu/__init__.py create mode 100644 zhenxun/plugins/send_setu_/send_setu/_data_source.py create mode 100644 zhenxun/plugins/send_setu_/update_setu/__init__.py create mode 100644 zhenxun/plugins/send_setu_/update_setu/data_source.py diff --git a/.vscode/settings.json b/.vscode/settings.json index 1f095bdd..ac493193 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -14,9 +14,11 @@ "hibiapi", "httpx", "kaiheila", + "lolicon", "nonebot", "onebot", "pixiv", + "Setu", "tobytes", "unban", "userinfo", diff --git a/poetry.lock b/poetry.lock index 35afa584..f1770cb4 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1684,6 +1684,85 @@ type = "legacy" url = "https://mirrors.aliyun.com/pypi/simple" reference = "ali" +[[package]] +name = "numpy" +version = "1.26.4" +description = "Fundamental package for array computing in Python" +optional = false +python-versions = ">=3.9" +files = [ + {file = "numpy-1.26.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9ff0f4f29c51e2803569d7a51c2304de5554655a60c5d776e35b4a41413830d0"}, + {file = "numpy-1.26.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2e4ee3380d6de9c9ec04745830fd9e2eccb3e6cf790d39d7b98ffd19b0dd754a"}, + {file = "numpy-1.26.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d209d8969599b27ad20994c8e41936ee0964e6da07478d6c35016bc386b66ad4"}, + {file = "numpy-1.26.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f"}, + {file = "numpy-1.26.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:62b8e4b1e28009ef2846b4c7852046736bab361f7aeadeb6a5b89ebec3c7055a"}, + {file = "numpy-1.26.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:a4abb4f9001ad2858e7ac189089c42178fcce737e4169dc61321660f1a96c7d2"}, + {file = "numpy-1.26.4-cp310-cp310-win32.whl", hash = "sha256:bfe25acf8b437eb2a8b2d49d443800a5f18508cd811fea3181723922a8a82b07"}, + {file = "numpy-1.26.4-cp310-cp310-win_amd64.whl", hash = "sha256:b97fe8060236edf3662adfc2c633f56a08ae30560c56310562cb4f95500022d5"}, + {file = "numpy-1.26.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4c66707fabe114439db9068ee468c26bbdf909cac0fb58686a42a24de1760c71"}, + {file = "numpy-1.26.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:edd8b5fe47dab091176d21bb6de568acdd906d1887a4584a15a9a96a1dca06ef"}, + {file = "numpy-1.26.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7ab55401287bfec946ced39700c053796e7cc0e3acbef09993a9ad2adba6ca6e"}, + {file = "numpy-1.26.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:666dbfb6ec68962c033a450943ded891bed2d54e6755e35e5835d63f4f6931d5"}, + {file = "numpy-1.26.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:96ff0b2ad353d8f990b63294c8986f1ec3cb19d749234014f4e7eb0112ceba5a"}, + {file = "numpy-1.26.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:60dedbb91afcbfdc9bc0b1f3f402804070deed7392c23eb7a7f07fa857868e8a"}, + {file = "numpy-1.26.4-cp311-cp311-win32.whl", hash = "sha256:1af303d6b2210eb850fcf03064d364652b7120803a0b872f5211f5234b399f20"}, + {file = "numpy-1.26.4-cp311-cp311-win_amd64.whl", hash = "sha256:cd25bcecc4974d09257ffcd1f098ee778f7834c3ad767fe5db785be9a4aa9cb2"}, + {file = "numpy-1.26.4-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b3ce300f3644fb06443ee2222c2201dd3a89ea6040541412b8fa189341847218"}, + {file = "numpy-1.26.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:03a8c78d01d9781b28a6989f6fa1bb2c4f2d51201cf99d3dd875df6fbd96b23b"}, + {file = "numpy-1.26.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9fad7dcb1aac3c7f0584a5a8133e3a43eeb2fe127f47e3632d43d677c66c102b"}, + {file = "numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed"}, + {file = "numpy-1.26.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:ab47dbe5cc8210f55aa58e4805fe224dac469cde56b9f731a4c098b91917159a"}, + {file = "numpy-1.26.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0"}, + {file = "numpy-1.26.4-cp312-cp312-win32.whl", hash = "sha256:50193e430acfc1346175fcbdaa28ffec49947a06918b7b92130744e81e640110"}, + {file = "numpy-1.26.4-cp312-cp312-win_amd64.whl", hash = "sha256:08beddf13648eb95f8d867350f6a018a4be2e5ad54c8d8caed89ebca558b2818"}, + {file = "numpy-1.26.4-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:7349ab0fa0c429c82442a27a9673fc802ffdb7c7775fad780226cb234965e53c"}, + {file = "numpy-1.26.4-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:52b8b60467cd7dd1e9ed082188b4e6bb35aa5cdd01777621a1658910745b90be"}, + {file = "numpy-1.26.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d5241e0a80d808d70546c697135da2c613f30e28251ff8307eb72ba696945764"}, + {file = "numpy-1.26.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f870204a840a60da0b12273ef34f7051e98c3b5961b61b0c2c1be6dfd64fbcd3"}, + {file = "numpy-1.26.4-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:679b0076f67ecc0138fd2ede3a8fd196dddc2ad3254069bcb9faf9a79b1cebcd"}, + {file = "numpy-1.26.4-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:47711010ad8555514b434df65f7d7b076bb8261df1ca9bb78f53d3b2db02e95c"}, + {file = "numpy-1.26.4-cp39-cp39-win32.whl", hash = "sha256:a354325ee03388678242a4d7ebcd08b5c727033fcff3b2f536aea978e15ee9e6"}, + {file = "numpy-1.26.4-cp39-cp39-win_amd64.whl", hash = "sha256:3373d5d70a5fe74a2c1bb6d2cfd9609ecf686d47a2d7b1d37a8f3b6bf6003aea"}, + {file = "numpy-1.26.4-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:afedb719a9dcfc7eaf2287b839d8198e06dcd4cb5d276a3df279231138e83d30"}, + {file = "numpy-1.26.4-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95a7476c59002f2f6c590b9b7b998306fba6a5aa646b1e22ddfeaf8f78c3a29c"}, + {file = "numpy-1.26.4-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:7e50d0a0cc3189f9cb0aeb3a6a6af18c16f59f004b866cd2be1c14b36134a4a0"}, + {file = "numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010"}, +] + +[package.source] +type = "legacy" +url = "https://mirrors.aliyun.com/pypi/simple" +reference = "ali" + +[[package]] +name = "opencv-python" +version = "4.9.0.80" +description = "Wrapper package for OpenCV python bindings." +optional = false +python-versions = ">=3.6" +files = [ + {file = "opencv-python-4.9.0.80.tar.gz", hash = "sha256:1a9f0e6267de3a1a1db0c54213d022c7c8b5b9ca4b580e80bdc58516c922c9e1"}, + {file = "opencv_python-4.9.0.80-cp37-abi3-macosx_10_16_x86_64.whl", hash = "sha256:7e5f7aa4486651a6ebfa8ed4b594b65bd2d2f41beeb4241a3e4b1b85acbbbadb"}, + {file = "opencv_python-4.9.0.80-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:71dfb9555ccccdd77305fc3dcca5897fbf0cf28b297c51ee55e079c065d812a3"}, + {file = "opencv_python-4.9.0.80-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b34a52e9da36dda8c151c6394aed602e4b17fa041df0b9f5b93ae10b0fcca2a"}, + {file = "opencv_python-4.9.0.80-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e4088cab82b66a3b37ffc452976b14a3c599269c247895ae9ceb4066d8188a57"}, + {file = "opencv_python-4.9.0.80-cp37-abi3-win32.whl", hash = "sha256:dcf000c36dd1651118a2462257e3a9e76db789a78432e1f303c7bac54f63ef6c"}, + {file = "opencv_python-4.9.0.80-cp37-abi3-win_amd64.whl", hash = "sha256:3f16f08e02b2a2da44259c7cc712e779eff1dd8b55fdb0323e8cab09548086c0"}, +] + +[package.dependencies] +numpy = [ + {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, + {version = ">=1.23.5", markers = "python_version >= \"3.11\" and python_version < \"3.12\""}, + {version = ">=1.21.4", markers = "python_version >= \"3.10\" and platform_system == \"Darwin\" and python_version < \"3.11\""}, + {version = ">=1.21.2", markers = "platform_system != \"Darwin\" and python_version >= \"3.10\" and python_version < \"3.11\""}, +] + +[package.source] +type = "legacy" +url = "https://mirrors.aliyun.com/pypi/simple" +reference = "ali" + [[package]] name = "pillow" version = "9.5.0" @@ -3205,4 +3284,4 @@ reference = "ali" [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "11beb90d388207c12255f2de15ad66f40ede82677ceb966a93bc31ebd97977f3" +content-hash = "76aa9b04323c716cda8d3e79a552d35c3f2d96eac39682c6c9c6b59291cbd398" diff --git a/pyproject.toml b/pyproject.toml index af263298..d218ba31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ beautifulsoup4 = "^4.12.3" lxml = "^5.1.0" psutil = "^5.9.8" feedparser = "^6.0.11" +opencv-python = "^4.9.0.80" [tool.poetry.dev-dependencies] diff --git a/zhenxun/plugins/send_setu_/__init__.py b/zhenxun/plugins/send_setu_/__init__.py new file mode 100644 index 00000000..eb35e275 --- /dev/null +++ b/zhenxun/plugins/send_setu_/__init__.py @@ -0,0 +1,5 @@ +from pathlib import Path + +import nonebot + +nonebot.load_plugins(str(Path(__file__).parent.resolve())) diff --git a/zhenxun/plugins/send_setu_/_model.py b/zhenxun/plugins/send_setu_/_model.py new file mode 100644 index 00000000..865af7d1 --- /dev/null +++ b/zhenxun/plugins/send_setu_/_model.py @@ -0,0 +1,87 @@ +from tortoise import fields +from tortoise.contrib.postgres.functions import Random +from tortoise.expressions import Q +from typing_extensions import Self + +from zhenxun.services.db_context import Model + + +class Setu(Model): + + id = fields.IntField(pk=True, generated=True, auto_increment=True) + """自增id""" + local_id = fields.IntField() + """本地存储下标""" + title = fields.CharField(255) + """标题""" + author = fields.CharField(255) + """作者""" + pid = fields.BigIntField() + """pid""" + img_hash = fields.TextField() + """图片hash""" + img_url = fields.CharField(255) + """pixiv url链接""" + is_r18 = fields.BooleanField() + """是否r18""" + tags = fields.TextField() + """tags""" + + class Meta: + table = "setu" + table_description = "色图数据表" + unique_together = ("pid", "img_url") + + @classmethod + async def query_image( + cls, + local_id: int | None = None, + tags: list[str] | None = None, + r18: bool = False, + limit: int = 50, + ) -> list[Self] | Self | None: + """通过tag查找色图 + + 参数: + local_id: 本地色图 id + tags: tags + r18: 是否 r18,0:非r18 1:r18 2:混合 + limit: 获取数量 + + 返回: + list[Self] | Self | None: 色图数据 + """ + if local_id: + return await cls.filter(is_r18=r18, local_id=local_id).first() + query = cls.filter(is_r18=r18) + if tags: + for tag in tags: + query = query.filter( + Q(tags__contains=tag) + | Q(title__contains=tag) + | Q(author__contains=tag) + ) + query = query.annotate(rand=Random()).limit(limit) + return await query.all() + + @classmethod + async def delete_image(cls, pid: int, img_url: str) -> int: + """删除图片并替换 + + 参数: + pid: 图片pid + + 返回: + int: 删除返回的本地id + """ + print(pid) + return_id = -1 + if query := await cls.get_or_none(pid=pid, img_url=img_url): + num = await cls.filter(is_r18=query.is_r18).count() + last_image = await cls.get_or_none(is_r18=query.is_r18, local_id=num - 1) + if last_image: + return_id = last_image.local_id + last_image.local_id = query.local_id + await last_image.save(update_fields=["local_id"]) + await query.delete() + return return_id diff --git a/zhenxun/plugins/send_setu_/send_setu/__init__.py b/zhenxun/plugins/send_setu_/send_setu/__init__.py new file mode 100644 index 00000000..94a82cb5 --- /dev/null +++ b/zhenxun/plugins/send_setu_/send_setu/__init__.py @@ -0,0 +1,227 @@ +import random +from typing import Tuple + +from nonebot.adapters import Bot +from nonebot.matcher import Matcher +from nonebot.message import run_postprocessor +from nonebot.plugin import PluginMetadata +from nonebot_plugin_alconna import ( + Alconna, + Args, + Arparma, + Match, + Option, + on_alconna, + store_true, +) +from nonebot_plugin_saa import Text +from nonebot_plugin_session import EventSession + +from zhenxun.configs.config import NICKNAME +from zhenxun.configs.utils import PluginCdBlock, PluginExtraData, RegisterConfig +from zhenxun.models.sign_user import SignUser +from zhenxun.models.user_console import UserConsole +from zhenxun.services.log import logger +from zhenxun.utils.withdraw_manage import WithdrawManager + +from ._data_source import SetuManage, base_config + +__plugin_meta__ = PluginMetadata( + name="色图", + description="不要小看涩图啊混蛋!", + usage=""" + 搜索 lolicon 图库,每日色图time... + 多个tag使用#连接 + 指令: + 色图: 随机色图 + 色图 -r: 随机在线r18涩图 + 色图 -id [id]: 本地指定id色图 + 色图 *[tags]: 在线搜索指定tag色图 + 色图 *[tags] -r: 同上, r18色图 + [1-9]张涩图: 本地随机色图连发 + [1-9]张[tags]的涩图: 在线搜索指定tag色图连发 + 示例:色图 萝莉|少女#白丝|黑丝 + 示例:色图 萝莉#猫娘 + 注: + tag至多取前20项,| 为或,萝莉|少女=萝莉或者少女 + """.strip(), + extra=PluginExtraData( + author="HibiKier", + version="0.1", + menu_type="来点好康的", + limits=[PluginCdBlock(result="您冲的太快了,请稍后再冲.")], + configs=[ + RegisterConfig( + key="WITHDRAW_SETU_MESSAGE", + value=(0, 1), + help="自动撤回,参1:延迟撤回色图时间(秒),0 为关闭 | 参2:监控聊天类型,0(私聊) 1(群聊) 2(群聊+私聊)", + default_value=(0, 1), + type=Tuple[int, int], + ), + RegisterConfig( + key="ONLY_USE_LOCAL_SETU", + value=False, + help="仅仅使用本地色图,不在线搜索", + default_value=False, + type=bool, + ), + RegisterConfig( + key="INITIAL_SETU_PROBABILITY", + value=0.7, + help="初始色图概率,总概率 = 初始色图概率 + 好感度", + default_value=0.7, + type=float, + ), + RegisterConfig( + key="DOWNLOAD_SETU", + value=True, + help="是否存储下载的色图,使用本地色图可以加快图片发送速度", + default_value=True, + type=float, + ), + RegisterConfig( + key="TIMEOUT", + value=10, + help="色图下载超时限制(秒)", + default_value=10, + type=int, + ), + RegisterConfig( + key="SHOW_INFO", + value=True, + help="是否显示色图的基本信息,如PID等", + default_value=True, + type=bool, + ), + RegisterConfig( + key="ALLOW_GROUP_R18", + value=False, + help="在群聊中启用R18权限", + default_value=False, + type=bool, + ), + RegisterConfig( + key="MAX_ONCE_NUM2FORWARD", + value=None, + help="单次发送的图片数量达到指定值时转发为合并消息", + default_value=None, + type=int, + ), + RegisterConfig( + key="MAX_ONCE_NUM", + value=10, + help="单次发送图片数量限制", + default_value=10, + type=int, + ), + RegisterConfig( + module="pixiv", + key="PIXIV_NGINX_URL", + value="i.pixiv.re", + help="Pixiv反向代理", + default_value="i.pixiv.re", + ), + ], + ).dict(), +) + + +@run_postprocessor +async def _( + matcher: Matcher, + exception: Exception | None, + session: EventSession, +): + if matcher.plugin_name == "send_setu": + # 添加数据至数据库 + try: + await SetuManage.save_to_database() + logger.info("色图数据自动存储数据库成功...") + except Exception: + pass + + +_matcher = on_alconna( + Alconna( + "色图", + Args["tags?", str], + Option("-n", Args["num", int, 1], help_text="数量"), + Option("-id", Args["local_id", int], help_text="本地id"), + Option("-r", action=store_true, help_text="r18"), + ), + aliases={"涩图", "不够色", "来一发", "再来点"}, + priority=5, + block=True, +) + +_matcher.shortcut( + r".*?(?P\d*)[份|发|张|个|次|点](?P.*)[瑟|色|涩]图.*?", + command="色图", + arguments=["{tags}", "-n", "{num}"], + prefix=True, +) + + +@_matcher.handle() +async def _( + bot: Bot, + session: EventSession, + arparma: Arparma, + num: Match[int], + tags: Match[str], + local_id: Match[int], +): + _tags = tags.result.split("#") if tags.available else None + if _tags and NICKNAME in _tags: + await Text( + "咳咳咳,虽然我很可爱,但是我木有自己的色图~~~有的话记得发我一份呀" + ).finish() + if not session.id1: + await Text("用户id为空...").finish() + gid = session.id3 or session.id2 + user_console = await UserConsole.get_user(session.id1, session.platform) + user, _ = await SignUser.get_or_create( + user_id=session.id1, + defaults={"user_console": user_console, "platform": session.platform}, + ) + if session.id1 not in bot.config.superusers: + """超级用户跳过罗翔""" + if result := SetuManage.get_luo(float(user.impression)): + await result.finish() + is_r18 = arparma.find("r") + _num = num.result + if is_r18 and gid: + """群聊中禁止查看r18""" + if not base_config.get("ALLOW_GROUP_R18"): + await Text( + random.choice( + [ + "这种不好意思的东西怎么可能给这么多人看啦", + "羞羞脸!给我滚出克私聊!", + "变态变态变态变态大变态!", + ] + ) + ).finish() + if local_id.available: + """指定id""" + result = await SetuManage.get_setu(local_id=local_id.result) + if isinstance(result, str): + await Text(result).finish(reply=True) + await result[0].finish() + result_list = await SetuManage.get_setu(tags=_tags, num=_num, is_r18=is_r18) + if isinstance(result_list, str): + await Text(result_list).finish(reply=True) + for result in result_list: + logger.info(f"发送色图 {result}", arparma.header_result, session=session) + receipt = await result.send() + if receipt: + message_id = receipt.extract_message_id().message_id # type: ignore + await WithdrawManager.withdraw_message( + bot, + message_id, + base_config.get("WITHDRAW_SETU_MESSAGE"), + session, + ) + logger.info( + f"调用发送 {num}张 色图 tags: {_tags}", arparma.header_result, session=session + ) diff --git a/zhenxun/plugins/send_setu_/send_setu/_data_source.py b/zhenxun/plugins/send_setu_/send_setu/_data_source.py new file mode 100644 index 00000000..796578b9 --- /dev/null +++ b/zhenxun/plugins/send_setu_/send_setu/_data_source.py @@ -0,0 +1,362 @@ +import os +import random +from pathlib import Path + +from asyncpg import UniqueViolationError +from nonebot_plugin_saa import Image, MessageFactory, Text +from pydantic import BaseModel + +from zhenxun.configs.config import NICKNAME, Config +from zhenxun.configs.path_config import IMAGE_PATH, TEMP_PATH +from zhenxun.services.log import logger +from zhenxun.utils.http_utils import AsyncHttpx +from zhenxun.utils.image_utils import compressed_image +from zhenxun.utils.utils import change_img_md5, change_pixiv_image_links + +from .._model import Setu + +headers = { + "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.6;" + " rv:2.0.1) Gecko/20100101 Firefox/4.0.1", + "Referer": "https://www.pixiv.net", +} + +base_config = Config.get("send_setu") + + +class SetuManage: + + URL = "https://api.lolicon.app/setu/v2" + save_data = [] + + @classmethod + async def get_setu( + cls, + *, + local_id: int | None = None, + num: int = 10, + tags: list[str] | None = None, + is_r18: bool = False, + ) -> list[MessageFactory] | str: + """获取色图 + + 参数: + local_id: 指定图片id + num: 数量 + tags: 标签 + is_r18: 是否r18 + + 返回: + list[MessageFactory] | str: 色图数据列表或消息 + + """ + result_list = [] + if local_id: + """本地id""" + data_list = await cls.get_setu_list(local_id=local_id) + if isinstance(data_list, str): + return data_list + file = await cls.get_image(data_list[0]) + if isinstance(file, str): + return file + return [cls.init_image_message(file, data_list[0])] + if base_config.get("ONLY_USE_LOCAL_SETU"): + """仅使用本地色图""" + flag = False + data_list = await cls.get_setu_list(tags=tags, is_r18=is_r18) + if isinstance(data_list, str): + return data_list + cls.save_data = data_list + if num > len(data_list): + num = len(data_list) + flag = True + setu_list = random.sample(data_list, num) + for setu in setu_list: + base_path = None + if setu.is_r18: + base_path = IMAGE_PATH / "_r18" + else: + base_path = IMAGE_PATH / "_setu" + file_path = base_path / f"{setu.local_id}.jpg" + if not file_path.exists(): + return f"本地色图Id: {setu.local_id} 不存在..." + result_list.append(cls.init_image_message(file_path, setu)) + if flag: + result_list.append( + MessageFactory([Text("坏了,已经没图了,被榨干了!")]) + ) + return result_list + data_list = await cls.search_lolicon(tags, num, is_r18) + if isinstance(data_list, str): + """搜索失败, 从本地数据库中搜索""" + data_list = await cls.get_setu_list(tags=tags, is_r18=is_r18) + if isinstance(data_list, str): + return data_list + if not data_list: + return "没找到符合条件的色图..." + cls.save_data = data_list + flag = False + if num > len(data_list): + num = len(data_list) + flag = True + for setu in data_list: + file = await cls.get_image(setu) + if isinstance(file, str): + result_list.append(MessageFactory([Text(file)])) + continue + result_list.append(cls.init_image_message(file, setu)) + if not result_list: + return "没找到符合条件的色图..." + if flag: + result_list.append(MessageFactory([Text("坏了,已经没图了,被榨干了!")])) + return result_list + + @classmethod + def init_image_message(cls, file: Path, setu: Setu) -> MessageFactory: + """初始化图片发送消息 + + 参数: + file: 图片路径 + setu: Setu + + 返回: + MessageFactory: 发送消息内容 + """ + data_list = [] + if base_config.get("SHOW_INFO"): + data_list.append( + Text( + f"id:{setu.local_id or ''}\n" + f"title:{setu.title}\n" + f"author:{setu.author}\n" + f"PID:{setu.pid}\n" + ) + ) + data_list.append(Image(file)) + return MessageFactory(data_list) + + @classmethod + async def get_setu_list( + cls, + *, + local_id: int | None = None, + tags: list[str] | None = None, + is_r18: bool = False, + ) -> list[Setu] | str: + """获取数据库中的色图数据 + + 参数: + local_id: 色图本地id. + tags: 标签. + is_r18: 是否r18. + + 返回: + list[Setu] | str: 色图数据列表或消息 + """ + image_list: list[Setu] = [] + if local_id: + image_count = await Setu.filter(is_r18=is_r18).count() - 1 + if local_id < 0 or local_id > image_count: + return f"超过当前上下限!({image_count})" + image_list = [await Setu.query_image(local_id, r18=is_r18)] # type: ignore + elif tags: + image_list = await Setu.query_image(tags=tags, r18=is_r18) # type: ignore + else: + image_list = await Setu.query_image(r18=is_r18) # type: ignore + if not image_list: + return "没找到符合条件的色图..." + return image_list + + @classmethod + def get_luo(cls, impression: float) -> MessageFactory | None: + """罗翔 + + 参数: + impression: 好感度 + + 返回: + MessageFactory | None: 返回数据 + """ + if initial_setu_probability := base_config.get("INITIAL_SETU_PROBABILITY"): + probability = float(impression) + initial_setu_probability * 100 + if probability < random.randint(1, 101): + return MessageFactory( + [ + Text("我为什么要给你发这个?"), + Image( + IMAGE_PATH + / "luoxiang" + / random.choice(os.listdir(IMAGE_PATH / "luoxiang")) + ), + Text(f"\n(快向{NICKNAME}签到提升好感度吧!)"), + ] + ) + return None + + @classmethod + async def get_image(cls, setu: Setu) -> str | Path: + """下载图片 + + 参数: + setu: Setu + + 返回: + str | Path: 图片路径或返回消息 + """ + url = change_pixiv_image_links(setu.img_url) + index = setu.local_id if setu.local_id else random.randint(1, 100000) + file_name = f"{index}_temp_setu.jpg" + base_path = TEMP_PATH + if setu.local_id: + """本地图片存在直接返回""" + file_name = f"{index}.jpg" + if setu.is_r18: + base_path = IMAGE_PATH / "_r18" + else: + base_path = IMAGE_PATH / "_setu" + local_file = base_path / file_name + if local_file.exists(): + return local_file + file = base_path / file_name + download_success = False + for i in range(3): + logger.debug(f"尝试在线下载第 {i+1} 次", "色图") + try: + if await AsyncHttpx.download_file( + url, + file, + timeout=base_config.get("TIMEOUT"), + ): + download_success = True + if setu.local_id is not None: + if ( + os.path.getsize(base_path / f"{index}.jpg") + > 1024 * 1024 * 1.5 + ): + compressed_image( + base_path / f"{index}.jpg", + ) + change_img_md5(file) + logger.info(f"下载 lolicon 图片 {url} 成功, id:{index}") + break + except TimeoutError as e: + logger.error(f"下载图片超时", "色图", e=e) + except Exception as e: + logger.error(f"下载图片错误", "色图", e=e) + return file if download_success else "图片被小怪兽恰掉啦..!QAQ" + + @classmethod + async def search_lolicon( + cls, tags: list[str] | None, num: int, is_r18: bool + ) -> list[Setu] | str: + """搜索lolicon色图 + + 参数: + tags: 标签 + num: 数量 + is_r18: 是否r18 + + 返回: + list[Setu] | str: 色图数据或返回消息 + """ + params = { + "r18": 1 if is_r18 else 0, # 添加r18参数 0为否,1为是,2为混合 + "tag": tags, # 若指定tag + "num": 20, # 一次返回的结果数量 + "size": ["original"], + } + for count in range(3): + logger.debug(f"尝试获取图片URL第 {count+1} 次", "色图") + try: + response = await AsyncHttpx.get( + cls.URL, timeout=base_config.get("TIMEOUT"), params=params + ) + if response.status_code == 200: + data = response.json() + if not data["error"]: + data = data["data"] + result_list = cls.__handle_data(data) + num = num if num < len(data) else len(data) + random_list = random.sample(result_list, num) + if not random_list: + return "没找到符合条件的色图..." + return random_list + else: + return "没找到符合条件的色图..." + except TimeoutError as e: + logger.error(f"获取图片URL超时", "色图", e=e) + except Exception as e: + logger.error(f"访问页面错误", "色图", e=e) + return "我网线被人拔了..QAQ" + + @classmethod + def __handle_data(cls, data: dict) -> list[Setu]: + """lolicon数据处理 + + 参数: + data: lolicon数据 + + 返回: + list[Setu]: 整理的数据 + """ + result_list = [] + for i in range(len(data)): + img_url = data[i]["urls"]["original"] + img_url = change_pixiv_image_links(img_url) + title = data[i]["title"] + author = data[i]["author"] + pid = data[i]["pid"] + tags = [] + for j in range(len(data[i]["tags"])): + tags.append(data[i]["tags"][j]) + # if command != "色图r": + # if "R-18" in tags: + # tags.remove("R-18") + setu = Setu( + title=title, + author=author, + pid=pid, + img_url=img_url, + tags=",".join(tags), + is_r18="R-18" in tags, + ) + result_list.append(setu) + return result_list + + @classmethod + async def save_to_database(cls): + """存储色图数据到数据库 + + 参数: + data_list: 色图数据列表 + """ + set_list = [] + exists_list = [] + for data in cls.save_data: + if f"{data.pid}:{data.img_url}" not in exists_list: + exists_list.append(f"{data.pid}:{data.img_url}") + set_list.append(data) + if set_list: + create_list = [] + _cnt = 0 + _r18_cnt = 0 + for setu in set_list: + try: + if not await Setu.exists(pid=setu.pid, img_url=setu.img_url): + idx = await Setu.filter(is_r18=setu.is_r18).count() + setu.local_id = idx + (_r18_cnt if setu.is_r18 else _cnt) + setu.img_hash = "" + if setu.is_r18: + _r18_cnt += 1 + else: + _cnt += 1 + create_list.append(setu) + except UniqueViolationError: + pass + cls.save_data = [] + if create_list: + try: + await Setu.bulk_create(create_list, 10) + logger.debug(f"成功保存 {len(create_list)} 条色图数据") + except Exception as e: + logger.error("存储色图数据错误...", e=e) diff --git a/zhenxun/plugins/send_setu_/update_setu/__init__.py b/zhenxun/plugins/send_setu_/update_setu/__init__.py new file mode 100644 index 00000000..fc3aa3ac --- /dev/null +++ b/zhenxun/plugins/send_setu_/update_setu/__init__.py @@ -0,0 +1,59 @@ +from nonebot.permission import SUPERUSER +from nonebot.plugin import PluginMetadata +from nonebot.rule import to_me +from nonebot_plugin_alconna import Alconna, Arparma, on_alconna +from nonebot_plugin_apscheduler import scheduler +from nonebot_plugin_saa import Text +from nonebot_plugin_session import EventSession + +from zhenxun.configs.config import Config +from zhenxun.configs.utils import BaseBlock, PluginExtraData +from zhenxun.services.log import logger +from zhenxun.utils.enum import PluginType + +from .data_source import update_setu_img + +__plugin_meta__ = PluginMetadata( + name="更新色图", + description="更新数据库内存在的色图", + usage=""" + 更新数据库内存在的色图 + 指令: + 更新色图 + """.strip(), + extra=PluginExtraData( + author="HibiKier", + version="0.1", + plugin_type=PluginType.SUPERUSER, + limits=[BaseBlock(result="色图正在更新...")], + ).dict(), +) + +_matcher = on_alconna( + Alconna("更新色图"), rule=to_me(), permission=SUPERUSER, priority=1, block=True +) + + +@_matcher.handle() +async def _(session: EventSession, arparma: Arparma): + if Config.get_config("send_setu", "DOWNLOAD_SETU"): + await Text("开始更新色图...").send(reply=True) + result = await update_setu_img(True) + if result: + await Text(result).send() + logger.info("更新色图", arparma.header_result, session=session) + else: + await Text("更新色图配置未开启...").send() + + +# 更新色图 +@scheduler.scheduled_job( + "cron", + hour=4, + minute=30, +) +async def _(): + if Config.get_config("send_setu", "DOWNLOAD_SETU"): + result = await update_setu_img() + if result: + logger.info(result, "自动更新色图") diff --git a/zhenxun/plugins/send_setu_/update_setu/data_source.py b/zhenxun/plugins/send_setu_/update_setu/data_source.py new file mode 100644 index 00000000..07d217d6 --- /dev/null +++ b/zhenxun/plugins/send_setu_/update_setu/data_source.py @@ -0,0 +1,187 @@ +import os +import shutil +from datetime import datetime + +import nonebot +import ujson as json +from asyncpg.exceptions import UniqueViolationError +from nonebot.drivers import Driver +from PIL import UnidentifiedImageError + +from zhenxun.configs.config import Config +from zhenxun.configs.path_config import IMAGE_PATH, TEMP_PATH, TEXT_PATH +from zhenxun.services.log import logger +from zhenxun.utils.http_utils import AsyncHttpx +from zhenxun.utils.image_utils import compressed_image +from zhenxun.utils.utils import change_pixiv_image_links + +from .._model import Setu + +driver: Driver = nonebot.get_driver() + +_path = IMAGE_PATH + + +# 替换旧色图数据,修复local_id一直是50的问题 +@driver.on_startup +async def update_old_setu_data(): + path = TEXT_PATH + setu_data_file = path / "setu_data.json" + r18_data_file = path / "r18_setu_data.json" + if setu_data_file.exists() or r18_data_file.exists(): + index = 0 + r18_index = 0 + count = 0 + fail_count = 0 + for file in [setu_data_file, r18_data_file]: + if file.exists(): + data = json.load(open(file, "r", encoding="utf8")) + for x in data: + if file == setu_data_file: + idx = index + if "R-18" in data[x]["tags"]: + data[x]["tags"].remove("R-18") + else: + idx = r18_index + img_url = ( + data[x]["img_url"].replace("i.pixiv.cat", "i.pximg.net") + if "i.pixiv.cat" in data[x]["img_url"] + else data[x]["img_url"] + ) + # idx = r18_index if 'R-18' in data[x]["tags"] else index + try: + if not await Setu.exists(pid=data[x]["pid"], url=img_url): + await Setu.create( + local_id=idx, + title=data[x]["title"], + author=data[x]["author"], + pid=data[x]["pid"], + img_hash=data[x]["img_hash"], + img_url=img_url, + is_r18="R-18" in data[x]["tags"], + tags=",".join(data[x]["tags"]), + ) + count += 1 + if "R-18" in data[x]["tags"]: + r18_index += 1 + else: + index += 1 + logger.info( + f'添加旧色图数据成功 PID:{data[x]["pid"]} index:{idx}....' + ) + except UniqueViolationError: + fail_count += 1 + logger.info( + f'添加旧色图数据失败,色图重复 PID:{data[x]["pid"]} index:{idx}....' + ) + file.unlink() + setu_url_path = path / "setu_url.json" + setu_r18_url_path = path / "setu_r18_url.json" + if setu_url_path.exists(): + setu_url_path.unlink() + if setu_r18_url_path.exists(): + setu_r18_url_path.unlink() + logger.info( + f"更新旧色图数据完成,成功更新数据:{count} 条,累计失败:{fail_count} 条" + ) + + +# 删除色图rar文件夹 +shutil.rmtree(IMAGE_PATH / "setu_rar", ignore_errors=True) +shutil.rmtree(IMAGE_PATH / "r18_rar", ignore_errors=True) +shutil.rmtree(IMAGE_PATH / "rar", ignore_errors=True) + +headers = { + "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.6;" + " rv:2.0.1) Gecko/20100101 Firefox/4.0.1", + "Referer": "https://www.pixiv.net", +} + + +async def update_setu_img(flag: bool = False) -> str | None: + """更新色图 + + 参数: + flag: 是否手动更新. + + 返回: + str | None: 更新信息 + """ + image_list = await Setu.all().order_by("local_id") + image_list.reverse() + _success = 0 + error_info = [] + error_type = [] + count = 0 + for image in image_list: + count += 1 + path = _path / "_r18" if image.is_r18 else _path / "_setu" + local_image = path / f"{image.local_id}.jpg" + path.mkdir(exist_ok=True, parents=True) + TEMP_PATH.mkdir(exist_ok=True, parents=True) + if not local_image.exists() or not image.img_hash: + temp_file = TEMP_PATH / f"{image.local_id}.jpg" + if temp_file.exists(): + temp_file.unlink() + url_ = change_pixiv_image_links(image.img_url) + try: + if not await AsyncHttpx.download_file( + url_, TEMP_PATH / f"{image.local_id}.jpg" + ): + continue + _success += 1 + try: + if ( + os.path.getsize( + TEMP_PATH / f"{image.local_id}.jpg", + ) + > 1024 * 1024 * 1.5 + ): + compressed_image( + TEMP_PATH / f"{image.local_id}.jpg", + path / f"{image.local_id}.jpg", + ) + else: + logger.info( + f"不需要压缩,移动图片{TEMP_PATH}/{image.local_id}.jpg " + f"--> /{path}/{image.local_id}.jpg" + ) + os.rename( + TEMP_PATH / f"{image.local_id}.jpg", + path / f"{image.local_id}.jpg", + ) + except FileNotFoundError: + logger.warning(f"文件 {image.local_id}.jpg 不存在,跳过...") + continue + # img_hash = str(get_img_hash(f"{path}/{image.local_id}.jpg")) + image.img_hash = "" + await image.save(update_fields=["img_hash"]) + # await Setu.update_setu_data(image.pid, img_hash=img_hash) + except UnidentifiedImageError: + # 图片已删除 + unlink = False + with open(local_image, "r") as f: + if "404 Not Found" in f.read(): + unlink = True + if unlink: + local_image.unlink() + max_num = await Setu.delete_image(image.pid, image.img_url) + if (path / f"{max_num}.jpg").exists(): + os.rename(path / f"{max_num}.jpg", local_image) + logger.warning(f"更新色图 PID:{image.pid} 404,已删除并替换") + except Exception as e: + _success -= 1 + logger.error(f"更新色图 {image.local_id}.jpg 错误 {type(e)}: {e}") + if type(e) not in error_type: + error_type.append(type(e)) + error_info.append( + f"更新色图 {image.local_id}.jpg 错误 {type(e)}: {e}" + ) + else: + logger.info(f"更新色图 {image.local_id}.jpg 已存在") + if _success or error_info or flag: + return ( + f'{str(datetime.now()).split(".")[0]} 更新 色图 完成,本地存在 {count} 张,实际更新 {_success} 张,以下为更新时未知错误:\n' + + "\n".join(error_info), + ) + return None diff --git a/zhenxun/services/db_context.py b/zhenxun/services/db_context.py index 34e0264c..dfd02925 100644 --- a/zhenxun/services/db_context.py +++ b/zhenxun/services/db_context.py @@ -51,39 +51,39 @@ async def init(): i_bind = bind if not i_bind: i_bind = f"{sql_name}://{user}:{password}@{address}:{port}/{database}" - # try: - await Tortoise.init( - db_url=i_bind, - modules={"models": MODELS}, - # timezone="Asia/Shanghai" - ) - logger.info(f"Database loaded successfully!") - # except Exception as e: - # raise Exception(f"数据库连接错误... {type(e)}: {e}") - if SCRIPT_METHOD: - db = Tortoise.get_connection("default") - logger.debug( - f"即将运行SCRIPT_METHOD方法, 合计 {len(SCRIPT_METHOD)} 个..." + try: + await Tortoise.init( + db_url=i_bind, + modules={"models": MODELS}, + # timezone="Asia/Shanghai" ) - sql_list = [] - for module, func in SCRIPT_METHOD: - try: - if is_coroutine_callable(func): - sql = await func() - else: - sql = func() - if sql: - sql_list += sql - except Exception as e: - logger.debug(f"{module} 执行SCRIPT_METHOD方法出错...", e=e) - for sql in sql_list: - logger.debug(f"执行SQL: {sql}") - try: - await db.execute_query_dict(sql) - # await TestSQL.raw(sql) - except Exception as e: - logger.debug(f"执行SQL: {sql} 错误...", e=e) - await Tortoise.generate_schemas() + if SCRIPT_METHOD: + db = Tortoise.get_connection("default") + logger.debug( + f"即将运行SCRIPT_METHOD方法, 合计 {len(SCRIPT_METHOD)} 个..." + ) + sql_list = [] + for module, func in SCRIPT_METHOD: + try: + if is_coroutine_callable(func): + sql = await func() + else: + sql = func() + if sql: + sql_list += sql + except Exception as e: + logger.debug(f"{module} 执行SCRIPT_METHOD方法出错...", e=e) + for sql in sql_list: + logger.debug(f"执行SQL: {sql}") + try: + await db.execute_query_dict(sql) + # await TestSQL.raw(sql) + except Exception as e: + logger.debug(f"执行SQL: {sql} 错误...", e=e) + await Tortoise.generate_schemas() + logger.info(f"Database loaded successfully!") + except Exception as e: + raise Exception(f"数据库连接错误...", e=e) async def disconnect(): diff --git a/zhenxun/utils/image_utils.py b/zhenxun/utils/image_utils.py index ddb08407..4dadb7ff 100644 --- a/zhenxun/utils/image_utils.py +++ b/zhenxun/utils/image_utils.py @@ -4,8 +4,11 @@ import re from pathlib import Path from typing import Awaitable, Callable +import cv2 from nonebot.utils import is_coroutine_callable +from zhenxun.configs.path_config import IMAGE_PATH + from ._build_image import BuildImage, ColorAlias from ._build_mat import BuildMat, MatType from ._image_template import ImageTemplate, RowStyle @@ -337,3 +340,27 @@ async def build_sort_image( curr_h += img.height + 10 curr_w += max([x.width for x in ig]) + 30 return A + + +def compressed_image( + in_file: str | Path, + out_file: str | Path | None = None, + ratio: float = 0.9, +): + """压缩图片 + + 参数: + in_file: 被压缩的文件路径 + out_file: 压缩后输出的文件路径 + ratio: 压缩率,宽高 * 压缩率 + """ + in_file = IMAGE_PATH / in_file if isinstance(in_file, str) else in_file + if out_file: + out_file = IMAGE_PATH / out_file if isinstance(out_file, str) else out_file + else: + out_file = in_file + h, w, d = cv2.imread(str(in_file.absolute())).shape + img = cv2.resize( + cv2.imread(str(in_file.absolute())), (int(w * ratio), int(h * ratio)) + ) + cv2.imwrite(str(out_file.absolute()), img) diff --git a/zhenxun/utils/withdraw_manage.py b/zhenxun/utils/withdraw_manage.py index 33c7b2b4..b5b8d176 100644 --- a/zhenxun/utils/withdraw_manage.py +++ b/zhenxun/utils/withdraw_manage.py @@ -7,6 +7,7 @@ from nonebot.adapters.kaiheila import Bot as KaiheilaBot from nonebot.adapters.onebot.v11 import Bot as v11Bot from nonebot.adapters.onebot.v12 import Bot as v12Bot from nonebot_plugin_session import EventSession +from ruamel.yaml.comments import CommentedSeq from zhenxun.services.log import logger @@ -80,7 +81,7 @@ class WithdrawManager: if time: gid = None _time = 1 - if isinstance(time, tuple): + if isinstance(time, (tuple, CommentedSeq)): if time[0] == 0: return if session: