zhenxun_bot/plugins/pix_gallery/pix.py

from nonebot.adapters.cqhttp.message import Message
from utils.utils import get_message_text, is_number
from configs.config import PIX_OMEGA_PIXIV_RATIO, WITHDRAW_PIX_TIME
from models.omega_pixiv_illusts import OmegaPixivIllusts
from utils.message_builder import image
from services.log import logger
from nonebot.adapters.cqhttp import (
    Bot,
    MessageEvent,
    GroupMessageEvent,
    PrivateMessageEvent,
)
from utils.static_data import withdraw_message_id_manager
from nonebot.typing import T_State
from .data_source import get_image
from models.pixiv import Pixiv
from nonebot import on_command
import random


pix = on_command("pix", aliases={"PIX", "Pix"}, priority=5, block=True)


PIX_RATIO = PIX_OMEGA_PIXIV_RATIO[0] / (
    PIX_OMEGA_PIXIV_RATIO[0] + PIX_OMEGA_PIXIV_RATIO[1]
)
OMEGA_RATIO = 1 - PIX_RATIO


@pix.handle()
async def _(bot: Bot, event: MessageEvent, state: T_State):
    num = 1
    keyword = get_message_text(event.json())
    if is_number(keyword):
        all_image = await Pixiv.query_images(uid=int(keyword))
    elif keyword.lower().startswith("pid"):
        pid = keyword.replace("pid", "").replace(":", "")
        if not is_number(pid):
            await pix.finish("PID必须是数字...", at_sender=True)
        isr18 = 0 if isinstance(event, GroupMessageEvent) else 2
        nsfw_tag = 0 if isinstance(event, GroupMessageEvent) else None
        all_image = await Pixiv.query_images(pid=int(pid), r18=isr18)
        if not all_image:
            all_image = await OmegaPixivIllusts.query_images(
                pid=int(pid), nsfw_tag=nsfw_tag
            )
    else:
        x = keyword.split()
        if '-s' in x:
            x.remove('-s')
            nsfw_tag = 1
        elif '-r' in x:
            x.remove('-r')
            nsfw_tag = 2
        else:
            nsfw_tag = 0
        if nsfw_tag != 0 and str(event.user_id) not in bot.config.superusers:
            await pix.finish('你不能看这些噢，这些都是是留给管理员看的...')
        if len(x) > 1:
            if is_number(x[-1]):
                num = int(x[-1])
                if num > 10:
                    if str(event.user_id) not in bot.config.superusers or (
                        str(event.user_id) in bot.config.superusers and num > 30
                    ):
                        num = random.randint(1, 10)
                        await pix.send(f"太贪心了，就给你发 {num}张 好了")
                x = x[:-1]
                keyword = " ".join(x)
        pix_num = int(num * PIX_RATIO) + 15 if PIX_RATIO != 0 else 0
        omega_num = num - pix_num + 15
        tmp = await Pixiv.query_images(
            x, r18=1 if nsfw_tag == 2 else 0, num=pix_num
        ) + await OmegaPixivIllusts.query_images(x, nsfw_tag=nsfw_tag, num=omega_num)
        tmp_ = []
        all_image = []
        for x in tmp:
            if x.pid not in tmp_:
                all_image.append(x)
                tmp_.append(x.pid)
    if not all_image:
        await pix.finish(f"未在图库中找到与 {keyword} 相关Tag/UID/PID的图片...", at_sender=True)
    for _ in range(num):
        img_url = None
        author = None
        if not all_image:
            await pix.finish("坏了...发完了，没图了...")
        img = random.choice(all_image)
        all_image.remove(img)
        if isinstance(img, OmegaPixivIllusts):
            img_url = img.url
            author = img.uname
            print(img.nsfw_tag)
        elif isinstance(img, Pixiv):
            img_url = img.img_url
            author = img.author
        pid = img.pid
        title = img.title
        uid = img.uid
        # tags = img.tags
        _img = await get_image(img_url, event.user_id)
        if _img:
            msg_id = await pix.send(
                Message(
                    f"title：{title}\n"
                    f"author：{author}\n"
                    f"PID：{pid}\nUID：{uid}\n"
                    f"{image(_img, 'temp')}"
                )
            )
            logger.info(
                f"(USER {event.user_id}, GROUP {event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
                f" 查看PIX图库PID: {pid}"
            )
        else:
            msg_id = await pix.send(f"下载图片似乎出了一点问题，PID：{pid}")
            logger.info(
                f"(USER {event.user_id}, GROUP {event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
                f" 查看PIX图库PID: {pid}，下载图片出错"
            )
        withdraw_message(event, msg_id["message_id"])


def withdraw_message(event: MessageEvent, id_: int):
    if WITHDRAW_PIX_TIME[0]:
        if (
            (WITHDRAW_PIX_TIME[1] == 0 and isinstance(event, PrivateMessageEvent))
            or (WITHDRAW_PIX_TIME[1] == 1 and isinstance(event, GroupMessageEvent))
            or WITHDRAW_PIX_TIME[1] == 2
        ):
            withdraw_message_id_manager["message_id"].append(
                (id_, WITHDRAW_PIX_TIME[0])
            )