diff options
| -rw-r--r-- | ATRI/plugins/essential.py | 1 | ||||
| -rw-r--r-- | ATRI/plugins/setu/__init__.py | 30 | ||||
| -rw-r--r-- | ATRI/plugins/setu/data_source.py | 8 | ||||
| -rw-r--r-- | ATRI/plugins/setu/tf_dealer.py | 4 |
4 files changed, 32 insertions, 11 deletions
diff --git a/ATRI/plugins/essential.py b/ATRI/plugins/essential.py index 4cc0ac2..09107d5 100644 --- a/ATRI/plugins/essential.py +++ b/ATRI/plugins/essential.py @@ -310,5 +310,6 @@ async def _recall_private_event(bot: Bot, event: FriendRecallNoticeEvent): async def _clear_cache(): try: shutil.rmtree(TEMP_PATH) + os.makedirs(TEMP_PATH, exist_ok=True) except Exception: log.warning("清除缓存失败，请手动清除：data/temp") diff --git a/ATRI/plugins/setu/__init__.py b/ATRI/plugins/setu/__init__.py index fb5736a..ffaae8d 100644 --- a/ATRI/plugins/setu/__init__.py +++ b/ATRI/plugins/setu/__init__.py @@ -77,6 +77,9 @@ async def _tag_setu(bot: Bot, event: MessageEvent): await bot.delete_msg(message_id=event_id) +_catcher_max_file_size = 128 + + setu_catcher = Setu().on_message("涩图嗅探", block=False) @@ -91,7 +94,7 @@ async def _setu_catcher(bot: Bot, event: MessageEvent): hso = list() for i in args: try: - data = await Setu().detecter(i) + data = await Setu().detecter(i, _catcher_max_file_size) except Exception: return if data[1] > 0.7: @@ -139,7 +142,7 @@ async def _deal_check(bot: Bot, event: MessageEvent, state: T_State): if not args: await nsfw_checker.reject("请发送图片而不是其他东西！！") - data = await Setu().detecter(args[0]) + data = await Setu().detecter(args[0], _catcher_max_file_size) hso = data[1] if not hso: await nsfw_checker.finish("图太小了！不测！") @@ -147,7 +150,7 @@ async def _deal_check(bot: Bot, event: MessageEvent, state: T_State): resu = f"涩值：{'{:.2%}'.format(hso)}\n" if hso >= 0.75: resu += "hso！不行我要发给别人看" - repo = f"涩图来咧！\n{MessageSegment.image(args[0])}\n涩值：{'{:.2%}'.format(hso[0])}" + repo = f"涩图来咧！\n{MessageSegment.image(args[0])}\n涩值：{'{:.2%}'.format(hso)}" for superuser in BotSelfConfig.superusers: await bot.send_private_msg(user_id=superuser, message=repo) @@ -159,6 +162,27 @@ async def _deal_check(bot: Bot, event: MessageEvent, state: T_State): await nsfw_checker.finish(resu) +catcher_setting = Setu().on_command("嗅探", "涩图检测图片文件大小设置") + +@catcher_setting.handle() +async def _catcher_setting(bot: Bot, event: MessageEvent, state: T_State): + msg = str(event.message).strip() + if msg: + state["catcher_set"] = msg + +@catcher_setting.got("catcher_set", "数值呢？（1对应1kb，默认128）") +async def _deal_setting(bot: Bot, event: MessageEvent, state: T_State): + global _catcher_max_file_size + msg = state["catcher_set"] + try: + _catcher_max_file_size = int(msg) + except Exception: + await catcher_setting.reject("请发送阿拉伯数字～！") + + repo = f"好诶！涩图检测文件最小值已设为：{_catcher_max_file_size}kb" + await catcher_setting.finish(repo) + + @scheduler.scheduled_job( "interval", name="涩批诱捕器", hours=1, misfire_grace_time=60, args=[Bot] ) diff --git a/ATRI/plugins/setu/data_source.py b/ATRI/plugins/setu/data_source.py index 39f3815..cca2767 100644 --- a/ATRI/plugins/setu/data_source.py +++ b/ATRI/plugins/setu/data_source.py @@ -62,11 +62,11 @@ class Setu(Service): return repo, setu @staticmethod - async def detecter(url: str) -> list: + async def detecter(url: str, file_size: int) -> list: """ 涩值检测. """ - data = await detect_image(url) + data = await detect_image(url, file_size) return data @staticmethod @@ -85,10 +85,6 @@ class Setu(Service): tag = choice(temp_data.get("tags", ["女孩子"])) - temp_arg = temp_data[0].get( - "urls", - "https://i.pixiv.cat/img-original/img/2021/02/28/22/44/49/88124144_p0.jpg", - ) url = temp_data[0]["urls"].get( "original", "https://i.pixiv.cat/img-original/img/2021/02/28/22/44/49/88124144_p0.jpg", diff --git a/ATRI/plugins/setu/tf_dealer.py b/ATRI/plugins/setu/tf_dealer.py index f41de49..bf68020 100644 --- a/ATRI/plugins/setu/tf_dealer.py +++ b/ATRI/plugins/setu/tf_dealer.py @@ -49,14 +49,14 @@ def prepare_image(img): return image -async def detect_image(url) -> list: +async def detect_image(url: str, file_size: int) -> list: try: req = await request.get(url) except RequestError: raise RequestError("Get info from download image failed!") img_byte = getsizeof(req.read()) // 1024 - if img_byte < 256: + if img_byte < file_size: return [0, 0] try: |