summaryrefslogtreecommitdiff
path: root/ATRI/plugins/setu/data_source.py
blob: 03f19a49dd0772e63cbf6332583be148fdac0a9b (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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import re
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 LoliconResponse, SetuInfo
from .nsfw_checker import detect_image, init_model


_LOLICON_URL = "https://api.lolicon.app/setu/v2"


class Setu:
    def __init__(self, url: str):
        self.url = url

    @classmethod
    async def new(cls, tag: str = str()) -> Tuple[MessageSegment, SetuInfo]:
        """new 一个涩图

        Args:
            tag (str, optional): 附加 tag, 默认无

        Raises:
            RequestError: 涩图请求失败

        Returns:
            Tuple[MessageSegment, dict]: 涩图本体, 涩图信息
        """
        url = _LOLICON_URL + (f"?tag={tag}" if tag else str())
        try:
            req = await request.get(url)
        except Exception:
            raise RequestError("setu: 请求失败")

        data = LoliconResponse.parse_obj(req.json()).data[0]
        title = data.title
        pid = data.pid
        url = data.urls.original

        if conf.Setu.reverse_proxy:
            patt = "://(.*?)/"
            domain = re.findall(patt, url)[0]
            setu = url.replace(domain, conf.Setu.reverse_proxy_domain)

        setu_data = SetuInfo(title=title, pid=pid, url=url)
        setu = MessageSegment.image(
            file=url,
            timeout=114514,
        )

        return setu, setu_data

    async def detecter(self, max_size: int, disab_gif: bool) -> float:
        """图片涩值检测

        Args:
            max_size (int): 检测文件大小限制
            disab_gif (bool): 是否检测动图

        Returns:
            float: 百分比涩值
        """
        return await detect_image(self.url, max_size, disab_gif)


from ATRI import driver

driver().on_startup(init_model)