summaryrefslogtreecommitdiff
path: root/ATRI/plugins/setu
diff options
context:
space:
mode:
authorKyomotoi <[email protected]>2022-02-03 14:36:24 +0800
committerKyomotoi <[email protected]>2022-02-03 14:36:24 +0800
commit3e32ca3964ff8f40e0b491e87f153040f2348fd0 (patch)
treed04d19ba5a25f6f1f6a9c4f8c398d49eb252df3c /ATRI/plugins/setu
parentf5a020d45f7294214bbcd488955b9c391d651a6d (diff)
downloadATRI-3e32ca3964ff8f40e0b491e87f153040f2348fd0.tar.gz
ATRI-3e32ca3964ff8f40e0b491e87f153040f2348fd0.tar.bz2
ATRI-3e32ca3964ff8f40e0b491e87f153040f2348fd0.zip
🔖 更新版本:
更新记录请参考文档: atri.kyomotoi.moe/changelog/overview/
Diffstat (limited to 'ATRI/plugins/setu')
-rw-r--r--ATRI/plugins/setu/__init__.py91
-rw-r--r--ATRI/plugins/setu/data_source.py23
-rw-r--r--ATRI/plugins/setu/tf_dealer.py10
3 files changed, 93 insertions, 31 deletions
diff --git a/ATRI/plugins/setu/__init__.py b/ATRI/plugins/setu/__init__.py
index 244e030..6c135aa 100644
--- a/ATRI/plugins/setu/__init__.py
+++ b/ATRI/plugins/setu/__init__.py
@@ -1,10 +1,11 @@
import re
import asyncio
from random import choice
+
from nonebot.permission import SUPERUSER
-from nonebot.adapters.cqhttp import Bot, MessageEvent, Message
-from nonebot.adapters.cqhttp.message import MessageSegment
-from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import CommandArg, ArgPlainText
+from nonebot.adapters.onebot.v11 import Bot, MessageEvent, Message, MessageSegment
from ATRI.config import BotSelfConfig
from ATRI.utils.limit import FreqLimiter, DailyLimiter
@@ -22,7 +23,7 @@ random_setu = Setu().on_command(
@random_setu.handle()
-async def _random_setu(bot: Bot, event: MessageEvent):
+async def _random_setu(bot: Bot, event: MessageEvent, matcher: Matcher, args: Message = CommandArg()):
user_id = event.get_user_id()
if not _setu_flmt.check(user_id):
await random_setu.finish()
@@ -44,12 +45,25 @@ async def _random_setu(bot: Bot, event: MessageEvent):
await asyncio.sleep(30)
await bot.delete_msg(message_id=event_id)
+ msg = args.extract_plain_text()
+ if msg:
+ matcher.set_arg("r_rush_after_think", args)
+
+@random_setu.got("r_rush_after_think")
+async def _(think: str = ArgPlainText("r_rush_after_think")):
+ is_repo = will_think(think)
+ if not is_repo:
+ await random_setu.finish()
+ else:
+ await random_setu.finish(is_repo)
+
+
tag_setu = Setu().on_regex(r"来[张点丶份](.*?)的[涩色🐍]图", "根据提供的tag查找涩图")
@tag_setu.handle()
-async def _tag_setu(bot: Bot, event: MessageEvent):
+async def _tag_setu(bot: Bot, event: MessageEvent, matcher: Matcher, args: Message = CommandArg()):
user_id = event.get_user_id()
if not _setu_flmt.check(user_id):
await random_setu.finish()
@@ -77,6 +91,19 @@ async def _tag_setu(bot: Bot, event: MessageEvent):
await asyncio.sleep(30)
await bot.delete_msg(message_id=event_id)
+ msg = args.extract_plain_text()
+ if msg:
+ matcher.set_arg("r_rush_after_think", args)
+
+
+@tag_setu.got("t_rush_after_think")
+async def _(think: str = ArgPlainText("t_rush_after_think")):
+ is_repo = will_think(think)
+ if not is_repo:
+ await random_setu.finish()
+ else:
+ await random_setu.finish(is_repo)
+
_catcher_max_file_size = 128
@@ -129,17 +156,16 @@ nsfw_checker = Setu().on_command("/nsfw", "涩值检测")
@nsfw_checker.handle()
-async def _nsfw_checker(bot: Bot, event: MessageEvent, state: T_State):
- msg = str(event.message).strip()
+async def _nsfw_checker(matcher: Matcher, args: Message = CommandArg()):
+ msg = args.extract_plain_text()
if msg:
- state["nsfw_img"] = msg
+ matcher.set_arg("nsfw_img", args)
@nsfw_checker.got("nsfw_img", "图呢?")
-async def _deal_check(bot: Bot, event: MessageEvent, state: T_State):
- msg = state["nsfw_img"]
+async def _deal_check(bot: Bot, img: str = ArgPlainText("nsfw_img")):
pattern = r"url=(.*?)]"
- args = re.findall(pattern, msg)
+ args = re.findall(pattern, img)
if not args:
await nsfw_checker.reject("请发送图片而不是其他东西!!")
@@ -167,16 +193,15 @@ catcher_setting = Setu().on_command("嗅探设置", "涩图检测图片文件大
@catcher_setting.handle()
-async def _catcher_setting(bot: Bot, event: MessageEvent, state: T_State):
- msg = str(event.message).strip()
+async def _catcher_setting(matcher: Matcher, args: Message = CommandArg()):
+ msg = args.extract_plain_text()
if msg:
- state["catcher_set"] = msg
+ matcher.set_arg("catcher_set", args)
@catcher_setting.got("catcher_set", "数值呢?(1对应1kb,默认128)")
-async def _deal_setting(bot: Bot, event: MessageEvent, state: T_State):
+async def _deal_setting(msg: str = ArgPlainText("catcher_set")):
global _catcher_max_file_size
- msg = state["catcher_set"]
try:
_catcher_max_file_size = int(msg)
except Exception:
@@ -205,3 +230,37 @@ async def _scheduler_setu(bot):
except Exception:
pass
+
+
+_ag_l = ["涩图来", "来点涩图", "来份涩图"]
+_ag_patt = r"来[张点丶份](.*?)的[涩色🐍]图"
+
+_nice_patt = r"[hH好][sS涩色][oO哦]|[嗯恩摁社蛇🐍射]了|(硬|石更)了|[牛🐂][牛🐂]要炸了|[炼恋]起来|开?导"
+_nope_patt = r"不够[涩色]|就这|这也[是叫算]|[??]"
+_again_patt = r"再来一张|不够"
+
+_nice_repo = ["w", "好诶!", "ohh", "(///w///)", "🥵", "我也"]
+_nope_repo = ["那你来发", "爱看不看", "你看不看吧", "看这种类型的涩图,是一件多么美妙的事情"]
+_again_repo = ["没了...", "自己找去"]
+
+
+def will_think(msg: str) -> str:
+ if msg in _ag_l:
+ return str()
+
+ ag_jud = re.findall(_ag_patt, msg)
+ if ag_jud:
+ return str()
+
+ nice_jud = re.findall(_nice_patt, msg)
+ nope_jud = re.findall(_nope_patt, msg)
+ again_jud = re.findall(_again_patt, msg)
+
+ if nice_jud:
+ return choice(_nice_repo)
+ elif nope_jud:
+ return choice(_nope_repo)
+ elif again_jud:
+ return choice(_again_repo)
+ else:
+ return str()
diff --git a/ATRI/plugins/setu/data_source.py b/ATRI/plugins/setu/data_source.py
index cca2767..a26bbef 100644
--- a/ATRI/plugins/setu/data_source.py
+++ b/ATRI/plugins/setu/data_source.py
@@ -1,16 +1,16 @@
-import base64
-
-# from pathlib import Path
from random import choice
-from nonebot.adapters.cqhttp import MessageSegment
+from nonebot.adapters.onebot.v11 import MessageSegment
from ATRI.service import Service
from ATRI.rule import is_in_service
from ATRI.utils import request
+from ATRI.config import Setu as ST
from .tf_dealer import detect_image
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"
+
class Setu(Service):
@@ -31,9 +31,9 @@ class Setu(Service):
data: dict = temp_data[0]
title = data.get("title", "木陰のねこ")
p_id = data.get("pid", 88124144)
- url = data["urls"].get("original", "ignore")
+ url: str = data["urls"].get("original", "ignore")
- setu = MessageSegment.image(url, timeout=114514)
+ setu = MessageSegment.image(use_proxy(url), timeout=114514)
repo = f"Title: {title}\nPid: {p_id}"
return repo, setu
@@ -55,7 +55,7 @@ class Setu(Service):
p_id = data.get("pid", 88124144)
url = data["urls"].get(
"original",
- "https://i.pixiv.cat/img-original/img/2021/02/28/22/44/49/88124144_p0.jpg",
+ use_proxy(DEFAULT_SETU),
)
setu = MessageSegment.image(url, timeout=114514)
repo = f"Title: {title}\nPid: {p_id}"
@@ -87,8 +87,15 @@ class Setu(Service):
url = temp_data[0]["urls"].get(
"original",
- "https://i.pixiv.cat/img-original/img/2021/02/28/22/44/49/88124144_p0.jpg",
+ use_proxy(DEFAULT_SETU),
)
setu = MessageSegment.image(url, timeout=114514)
repo = f"是{tag}哦~❤\n{setu}"
return repo
+
+
+def use_proxy(url: str) -> str:
+ if ST.reverse_proxy:
+ return url.replace("i.pixiv.cat", ST.reverse_proxy_domain)
+ else:
+ return url
diff --git a/ATRI/plugins/setu/tf_dealer.py b/ATRI/plugins/setu/tf_dealer.py
index bf68020..f966636 100644
--- a/ATRI/plugins/setu/tf_dealer.py
+++ b/ATRI/plugins/setu/tf_dealer.py
@@ -1,7 +1,6 @@
import io
import os
import re
-import string
import asyncio
import skimage
import skimage.io
@@ -9,12 +8,8 @@ import numpy as np
from PIL import Image
from pathlib import Path
from sys import getsizeof
-from random import sample
-try:
- import tflite_runtime.interpreter as tf # type: ignore
-except Exception:
- import tensorflow as tf
+import tensorflow as tf
from ATRI.log import logger as log
from ATRI.utils import request
@@ -116,4 +111,5 @@ async def init_module():
raise WriteError("NSFW TF module init failed!")
-asyncio.get_event_loop().run_until_complete(init_module())
+loop = asyncio.get_event_loop()
+loop.create_task(init_module())