From 1ee2eae499b79bf0d22212184bffd57d873826cc Mon Sep 17 00:00:00 2001 From: Kyomotoi <0w0@imki.moe> Date: Fri, 2 Dec 2022 08:52:19 +0800 Subject: =?UTF-8?q?=F0=9F=8E=A8=20=E9=87=8D=E6=9E=84=E6=B6=A9=E5=9B=BE?= =?UTF-8?q?=E7=BB=84=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ATRI/plugins/setu/__init__.py | 27 +++++----- ATRI/plugins/setu/data_source.py | 110 ++++++++++++++++++-------------------- ATRI/plugins/setu/models.py | 6 +++ ATRI/plugins/setu/nsfw_checker.py | 4 +- 4 files changed, 74 insertions(+), 73 deletions(-) create mode 100644 ATRI/plugins/setu/models.py diff --git a/ATRI/plugins/setu/__init__.py b/ATRI/plugins/setu/__init__.py index b7dc5f4..4a9c6a6 100644 --- a/ATRI/plugins/setu/__init__.py +++ b/ATRI/plugins/setu/__init__.py @@ -29,13 +29,14 @@ async def _(): async def _random_setu(bot: Bot, event: MessageEvent): loop = asyncio.get_running_loop() - repo, se = await Setu.random_setu() - await bot.send(event, repo) + setu, setu_data = await Setu.new() + setu_info = f"Title: {setu_data.title}\nPid: {setu_data.pid}" + await bot.send(event, setu_info) try: - msg_1 = await bot.send(event, Message(se)) + msg_1 = await bot.send(event, setu) except Exception: - await random_setu.finish("hso（发不出") + await random_setu.finish("hso (发不出") msg_id = msg_1["message_id"] loop.call_later( @@ -56,23 +57,21 @@ async def _(think: str = ArgPlainText("r_rush_after_think")): tag_setu = plugin.on_regex(r"来[张点丶份](.*?)的[涩色🐍]图", "根据提供的tag查找涩图，冷却2分钟") -@tag_setu.handle([Cooldown(120, prompt="慢...慢一..点❤")]) +@tag_setu.handle([Cooldown(120, prompt="")]) async def _tag_setu(bot: Bot, event: MessageEvent): loop = asyncio.get_running_loop() msg = str(event.get_message()).strip() pattern = r"来[张点丶份](.*?)的[涩色🐍]图" tag = re.findall(pattern, msg)[0] - repo, se = await Setu.tag_setu(tag) - if not plugin: - await tag_setu.finish(repo) - - await bot.send(event, repo) + setu, setu_data = await Setu.new(tag) + setu_info = f"Title: {setu_data.title}\nPid: {setu_data.pid}" + await bot.send(event, setu_info) try: - msg_1 = await bot.send(event, Message(se)) + msg_1 = await bot.send(event, setu) except Exception: - await random_setu.finish("hso（发不出") + await random_setu.finish("hso (发不出") msg_id = msg_1["message_id"] loop.call_later( @@ -105,7 +104,7 @@ async def _setu_catcher(bot: Bot, event: MessageEvent): hso = list() for i in args: try: - data = await Setu.detecter(i, _catcher_max_file_size) + data = await Setu(i).detecter(_catcher_max_file_size) except Exception: return if data > 0.7: @@ -144,7 +143,7 @@ async def _deal_check(bot: Bot, event: MessageEvent): if not args: await nsfw_checker.reject("请发送图片而不是其他东西！！") - hso = await Setu.detecter(args[0], _catcher_max_file_size) + hso = await Setu(args[0]).detecter(_catcher_max_file_size) if not hso: await nsfw_checker.finish("图太小了！不测！") diff --git a/ATRI/plugins/setu/data_source.py b/ATRI/plugins/setu/data_source.py index 40cddd1..8a1a26d 100644 --- a/ATRI/plugins/setu/data_source.py +++ b/ATRI/plugins/setu/data_source.py @@ -1,76 +1,72 @@ -import asyncio -from nonebot.adapters.onebot.v11 import Bot, MessageSegment +from typing import Tuple +from nonebot.adapters.onebot.v11 import MessageSegment from ATRI import conf from ATRI.utils import request +from ATRI.exceptions import RequestError + +from .models import SetuInfo from .nsfw_checker import detect_image, init_model -LOLICON_URL = "https://api.lolicon.app/setu/v2" -DEFAULT_SETU = ( - "https://i.pixiv.cat/img-original/img/2021/02/28/22/44/49/88124144_p0.jpg" -) +__LOLICON_URL = "https://api.lolicon.app/setu/v2" class Setu: - @staticmethod - def _use_proxy(url: str) -> str: - if conf.Setu.reverse_proxy: - return url.replace("i.pixiv.cat", conf.Setu.reverse_proxy_domain) - else: - return url + def __init__(self, url: str): + self.url = url @classmethod - async def random_setu(cls) -> tuple: - """ - 随机涩图. - """ - res = await request.get(LOLICON_URL) - data: dict = res.json() - temp_data: dict = data.get("data", list()) - if not temp_data: - return "涩批爬", None + async def new(cls, tag: str = str()) -> Tuple[MessageSegment, SetuInfo]: + """new 一个涩图 - data: dict = temp_data[0] - title = data.get("title", "木陰のねこ") - p_id = data.get("pid", 88124144) - url: str = data["urls"].get("original", "ignore") + Args: + tag (str, optional): 附加 tag, 默认无 - setu = MessageSegment.image(cls._use_proxy(url), timeout=114514) - repo = f"Title: {title}\nPid: {p_id}" - return repo, setu + Raises: + RequestError: 涩图请求失败 - @classmethod - async def tag_setu(cls, tag: str) -> tuple: - """ - 指定tag涩图. + Returns: + Tuple[MessageSegment, dict]: 涩图本体, 涩图信息 """ - url = LOLICON_URL + f"?tag={tag}" - res = await request.get(url) - data: dict = res.json() - - temp_data: dict = data.get("data", list()) - if not temp_data: - return f"没有 {tag} 的涩图呢...", None - - data = temp_data[0] - title = data.get("title", "木陰のねこ") - p_id = data.get("pid", 88124144) - url = data["urls"].get( - "original", - cls._use_proxy(DEFAULT_SETU), - ) - setu = MessageSegment.image(url, timeout=114514) - repo = f"Title: {title}\nPid: {p_id}" - return repo, setu - - @staticmethod - async def detecter(url: str, file_size: int) -> float: - """ - 涩值检测. + url = __LOLICON_URL + if tag: + url = __LOLICON_URL + f"?tag={tag}" + try: + req = await request.get(url) + except Exception: + raise RequestError("setu: 请求失败") + + data = req.json() + cache_data = data.get("data") + if not cache_data: + raise RequestError("今天不可以涩") + + data = cache_data[0] + title = data["title"] + pid = data["pid"] + setu = data["urls"].get("original", "ignore") + + if conf.Setu.reverse_proxy: + setu = MessageSegment.image( + file=setu.replace("i.pixiv.cat", conf.Setu.reverse_proxy_domain), + timeout=114514, + ) + + setu_data = SetuInfo(title=title, pid=pid) + + return setu, setu_data + + async def detecter(self, max_size: int) -> float: + """图片涩值检测 + + Args: + max_size (int): 检测文件大小限制 + + Returns: + float: 百分比涩值 """ - data = await detect_image(url, file_size) - return data + return await detect_image(self.url, max_size) from ATRI import driver diff --git a/ATRI/plugins/setu/models.py b/ATRI/plugins/setu/models.py new file mode 100644 index 0000000..7144f27 --- /dev/null +++ b/ATRI/plugins/setu/models.py @@ -0,0 +1,6 @@ +from pydantic import BaseModel + + +class SetuInfo(BaseModel): + title: str + pid: str diff --git a/ATRI/plugins/setu/nsfw_checker.py b/ATRI/plugins/setu/nsfw_checker.py index c0bd2ba..53546f6 100644 --- a/ATRI/plugins/setu/nsfw_checker.py +++ b/ATRI/plugins/setu/nsfw_checker.py @@ -40,14 +40,14 @@ def prepare_image(img): return image -async def detect_image(url: str, file_size: int) -> float: +async def detect_image(url: str, max_size: int) -> float: try: req = await request.get(url) except Exception: raise RequestError("Get info from download image failed!") img_byte = getsizeof(req.read()) // 1024 - if img_byte < file_size: + if img_byte < max_size: return 0 try: -- cgit v1.2.3