summaryrefslogtreecommitdiff
path: root/ATRI/plugins/saucenao/__init__.py
blob: 092db50a26613ac80f965c9a1e62811e40cf8895 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from re import findall
from random import choice

from nonebot.matcher import Matcher
from nonebot.params import ArgPlainText, CommandArg
from nonebot.adapters.onebot.v11 import MessageEvent, Message, MessageSegment

from ATRI.config import SauceNAO
from ATRI.utils.limit import FreqLimiter
from .data_source import SaouceNao


_search_flmt = FreqLimiter(10)
_search_flmt_notice = choice(["慢...慢一..点❤", "冷静1下", "歇会歇会~~"])


saucenao = SaouceNao().on_command("以图搜图", "透过一张图搜索可能的来源")


@saucenao.handle()
async def _ready_search(
    matcher: Matcher, event: MessageEvent, args: Message = CommandArg()
):
    user_id = event.get_user_id()
    if not _search_flmt.check(user_id):
        await saucenao.finish(_search_flmt_notice)

    msg = args.extract_plain_text()
    if msg:
        matcher.set_arg("saucenao_img", args)


@saucenao.got("saucenao_img", "图呢?")
async def _deal_search(event: MessageEvent, msg: str = ArgPlainText("saucenao_img")):
    user_id = event.get_user_id()
    img = findall(r"url=(.*?)]", msg)
    if not img:
        await saucenao.reject("请发送图片而不是其他东西!!")

    try:
        a = SaouceNao(SauceNAO.key)
    except Exception:
        await saucenao.finish("失败了...")

    result = f"> {MessageSegment.at(user_id)}" + await a.search(img[0])  # type: ignore
    _search_flmt.start_cd(user_id)
    await saucenao.finish(Message(result))