summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--ATRI/plugins/setu/__init__.py27
-rw-r--r--ATRI/plugins/setu/data_source.py110
-rw-r--r--ATRI/plugins/setu/models.py6
-rw-r--r--ATRI/plugins/setu/nsfw_checker.py4
4 files changed, 74 insertions, 73 deletions
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:
Powered by cgit, Theme cgit-dark created by Ted Yin, updated by Mole Shang.