Lanly109 / headimg_generator

基于HoshinoBot的制作头像相关的表情包插件,移植自nonebot-plugin-petpet
MIT License
50 stars 11 forks source link

渣渣的我又来了,加了一个每日次数限制 #14

Closed jumpad closed 1 year ago

jumpad commented 1 year ago

纯小白不会弄分支,直接贴请见谅。 文件: __init__.py

import shlex
import traceback
import base64
from typing import List
from .utils import help_image
from hoshino.typing import CQEvent, MessageSegment
from nonebot import on_startup
from hoshino import Service, priv, HoshinoBot
from hoshino.util import DailyNumberLimiter
from .data_source import make_image, commands
from .download import DownloadError, ResourceError
from .models import UserInfo, Command

sv = Service('头像表情包', help_='''
![](https://s2.loli.net/2022/09/18/zrhvqdCMDyOSpZj.jpg)
'''.strip(), enable_on_default=True, bundle='娱乐', visible = True)
_lmt = DailyNumberLimiter(12)

@sv.on_fullmatch(("帮助头像表情包"))
async def help(bot, ev: CQEvent):
    img = await help_image(commands)
    base64_str = base64.b64encode(img.getvalue()).decode()
    img =  'base64://' + base64_str
    img = str(MessageSegment.image(img))
    await bot.send(ev, img)

def is_qq(msg: str):
    return msg.isdigit() and 11 >= len(msg) >= 5

async def get_user_info(bot: HoshinoBot, user: UserInfo, ev: CQEvent):
    uid = ev.user_id
    if not _lmt.check(uid):
        await bot.finish(ev, '已经玩6次了!今天你不许玩这个了!', at_sender=True)
    if not user.qq:
        return

    if user.group:
        info = await bot.get_group_member_info(
            group_id=int(user.group), user_id=int(user.qq)
        )
        user.name = info.get("card", "") or info.get("nickname", "")
        user.gender = info.get("sex", "")
    else:
        info = await bot.get_stranger_info(user_id=int(user.qq))
        user.name = info.get("nickname", "")
        user.gender = info.get("sex", "")
    _lmt.increase(uid)

class Handler:
    def __init__(self, command):
        self.command = command

    async def handle(self, bot, ev: CQEvent):
        users: List[UserInfo] = []
        args: List[str] = []
        msg = ev.message

        # if ev.reply:
        #     reply_imgs = ev.reply.message["image"]
        #     for reply_img in reply_imgs:
        #         users.append(UserInfo(img_url=reply_img.data["url"]))

        for msg_seg in msg:
            if msg_seg.type == "at":
                users.append(
                    UserInfo(
                        qq=msg_seg.data["qq"],
                        group=str(ev.group_id)
                    )
                )
            elif msg_seg.type == "image":
                users.append(UserInfo(img_url=msg_seg.data["url"]))
            elif msg_seg.type == "text":
                raw_text = str(msg_seg)
                try:
                    texts = shlex.split(raw_text)
                except:
                    texts = raw_text.split()
                for text in texts:
                    if is_qq(text):
                        users.append(UserInfo(qq=text))
                    elif text == "自己":
                        users.append(
                            UserInfo(
                                qq=str(ev.user_id),
                                group=str(ev.group_id)
                            )
                        )
                    else:
                        text = text.strip()
                        if text:
                            args.append(text)

        if len(args) > self.command.arg_num:
            return False
        if ev.raw_message.find(f"[CQ:at,qq={str(ev.self_id)}") != -1:
            users.append(UserInfo(qq=str(ev.self_id), group=str(ev.group_id)))
        if not users:
            return False

        sender = UserInfo(qq=str(ev.user_id))
        await get_user_info(bot, sender, ev)

        for user in users:
            await get_user_info(bot, user, ev)

        try:
            img = await make_image(command = self.command, sender = sender, users = users, args = args)
            base64_str = base64.b64encode(img.getvalue()).decode()
            img =  'base64://' + base64_str
            img = str(MessageSegment.image(img))
        except Exception as e:
            print(traceback.format_exc())
            img = str(e)

        await bot.send(ev, img)

@on_startup
async def register_handler():
    for command in commands:
        func = getattr(sv, "on_prefix")
        key = command.keywords
        func = func(key, only_to_me=False)
        func(Handler(command).handle)
    print('petpet register done.')