summaryrefslogtreecommitdiff
path: root/ATRI/plugins/setu/modules/store.py
blob: e278c7786cbd34621369369a6d166dfbe7170633 (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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import json
from random import choice

from nonebot.typing import T_State
from nonebot.permission import SUPERUSER
from nonebot.adapters.cqhttp import Bot, MessageEvent
from nonebot.adapters.cqhttp.message import Message, MessageSegment

from ATRI.service import Service as sv
from ATRI.utils.request import get_bytes
from ATRI.exceptions import RequestError

from .data_source import SetuData


API_URL: str = "https://api.kyomotoi.moe/api/pixiv/illust?id="


__doc__ = """
为本地添加涩图!
权限组:维护者
用法:
  添加涩图 (pid)
补充:
  pid: Pixiv 作品id
"""


add_setu = sv.on_command(cmd="添加涩图", docs=__doc__, permission=SUPERUSER)


@add_setu.args_parser  # type: ignore
async def _load_add_setu(bot: Bot, event: MessageEvent, state: T_State) -> None:
    msg = str(event.message).strip()
    cancel = ["算了", "罢了"]
    if msg in cancel:
        await add_setu.finish("好吧...")
    if not msg:
        await add_setu.reject("涩图(pid)速发!")
    else:
        state["setu_add"] = msg


@add_setu.handle()
async def _add_setu(bot: Bot, event: MessageEvent, state: T_State) -> None:
    msg = str(event.message).strip()
    if msg:
        state["setu_add"] = msg


@add_setu.got("setu_add", prompt="涩图(pid)速发!")
async def _deal_add_setu(bot: Bot, event: MessageEvent, state: T_State) -> None:
    pid = state["setu_add"]

    URL = API_URL + pid
    try:
        data = json.loads(await get_bytes(URL))["illust"]
    except RequestError:
        raise RequestError("Request failed!")

    try:
        pic = data["meta_single_page"]["original_image_url"].replace(
            "pximg.net", "pixiv.cat"
        )
    except Exception:
        pic = choice(data["meta_pages"])["image_urls"]["original"].replace(
            "pximg.net", "pixiv.cat"
        )

    d = {
        "pid": pid,
        "title": data["title"],
        "tags": str(data["tags"]),
        "user_id": data["user"]["id"],
        "user_name": data["user"]["name"],
        "user_account": data["user"]["account"],
        "url": pic,
    }
    await SetuData.add_data(d)

    show_img = data["image_urls"]["medium"].replace("pximg.net", "pixiv.cat")
    msg = (
        "好欸!是新涩图:\n"
        f"Pid: {pid}\n"
        f"Title: {data['title']}\n"
        f"{MessageSegment.image(show_img)}"
    )
    await add_setu.finish(Message(msg))


__doc__ = """
删除涩图!
权限组:维护者
用法:
  删除涩图 (pid)
补充:
  pid: Pixiv 作品id
"""


del_setu = sv.on_command(cmd="删除涩图", docs=__doc__, permission=SUPERUSER)


@del_setu.args_parser  # type: ignore
async def _load_del_setu(bot: Bot, event: MessageEvent, state: T_State) -> None:
    msg = str(event.message).strip()
    cancel = ["算了", "罢了"]
    if msg in cancel:
        await add_setu.finish("好吧...")
    if not msg:
        await add_setu.reject("涩图(pid)速发!")
    else:
        state["setu_del"] = msg


@del_setu.handle()
async def _del_setu(bot: Bot, event: MessageEvent, state: T_State) -> None:
    msg = str(event.message).strip()
    if msg:
        state["setu_del"] = msg


@del_setu.got("setu_del", prompt="涩图(pid)速发!")
async def _deal_del_setu(bot: Bot, event: MessageEvent, state: T_State) -> None:
    pid = int(state["setu_del"])
    await SetuData.del_data(pid)
    await del_setu.finish(f"涩图({pid})已删除...")


count_setu = sv.on_command(cmd="涩图总量", permission=SUPERUSER)


@count_setu.handle()
async def _count_setu(bot: Bot, event: MessageEvent) -> None:
    msg = f"咱本地搭载了 {await SetuData.count()} 张涩图!"
    await count_setu.finish(msg)