iShotaBot/modules/sticker_download.py

import contextlib
import tempfile
from typing import TYPE_CHECKING

from cashews import cache
from pyrogram import filters, ContinuePropagation
from pyrogram.enums import ChatAction, MessageEntityType
from pyrogram.errors import StickersetInvalid

from defs.sticker_download import get_from_sticker_set, get_from_sticker, custom_emoji_filter, \
    get_from_custom_emoji, export_end, export_add
from init import bot

if TYPE_CHECKING:
    from pyrogram import Client
    from pyrogram.types import Message


@bot.on_message(
    filters.private & filters.text & filters.incoming & filters.regex(r"^https://t.me/addstickers/.*")
)
async def process_sticker_set(client: "Client", message: "Message"):
    cid = message.from_user.id
    short_name = message.text.replace("https://t.me/addstickers/", "")
    file = None
    try:
        reply = await message.reply("正在下载贴纸包...请耐心等待", quote=True)
        file = await get_from_sticker_set(short_name, cid, client, reply)
        with contextlib.suppress(Exception):
            await message.reply_chat_action(ChatAction.UPLOAD_DOCUMENT)
        await message.reply_document(file.as_posix(), quote=True)
        with contextlib.suppress(Exception):
            await reply.delete()
    except StickersetInvalid:
        await message.reply("无效的贴纸包", quote=True)
    finally:
        if file:
            file.unlink(missing_ok=True)
    raise ContinuePropagation


@bot.on_message(filters.private & filters.sticker & filters.incoming)
async def process_single_sticker(client: "Client", message: "Message"):
    await message.reply_chat_action(ChatAction.TYPING)
    if temp := await cache.get(f"sticker:export:{message.from_user.id}"):
        await export_add(temp, message.sticker, client)
        await message.reply_text("成功加入导出列表，结束选择请输入 /sticker_export_end", quote=True)
    else:
        reply = await message.reply("正在转换贴纸...请耐心等待", quote=True)
        target_file = None
        try:
            target_file = await get_from_sticker(client, message)
            await message.reply_chat_action(ChatAction.UPLOAD_DOCUMENT)
            await message.reply_document(target_file.as_posix(), quote=True)
        finally:
            if target_file:
                target_file.unlink(missing_ok=True)
        with contextlib.suppress(Exception):
            await reply.delete()
    raise ContinuePropagation


@bot.on_message(filters.private & custom_emoji_filter & filters.incoming)
async def process_custom_emoji(client: "Client", message: "Message"):
    try:
        stickers = await client.get_custom_emoji_stickers(
            [i.custom_emoji_id for i in message.entities if i and i.type == MessageEntityType.CUSTOM_EMOJI]
        )
    except AttributeError:
        await message.reply("无法获取贴纸", quote=True)
        raise ContinuePropagation
    reply = await message.reply(f"正在下载 {len(stickers)} 个 emoji ...请耐心等待", quote=True)
    for sticker in stickers:
        target_file = None
        try:
            target_file = await get_from_custom_emoji(client, sticker)
            await message.reply_chat_action(ChatAction.UPLOAD_DOCUMENT)
            await message.reply_document(target_file.as_posix(), quote=True)
        finally:
            if target_file:
                target_file.unlink(missing_ok=True)
    with contextlib.suppress(Exception):
        await reply.delete()
    raise ContinuePropagation


@bot.on_message(
    filters.private & filters.incoming & filters.command(
        ["sticker_export_start", "sticker_export_end"]
    )
)
async def batch_start(_: "Client", message: "Message"):
    uid = message.from_user.id
    if "start" in message.command[0].lower():
        if await cache.get(f"sticker:export:{uid}"):
            await message.reply("已经开始批量导出贴纸，请直接发送贴纸，完成选择请输入 /sticker_export_end", quote=True)
            return
        await cache.set(f"sticker:export:{uid}", tempfile.mkdtemp())
        await message.reply("开始批量导出贴纸，请直接发送贴纸，完成选择请输入 /sticker_export_end", quote=True)
    else:
        target_dir = await cache.get(f"sticker:export:{uid}")
        if not target_dir:
            await message.reply("未开始批量导出贴纸，请先使用命令 /sticker_export_start", quote=True)
            return
        file = None
        try:
            reply = await message.reply("正在打包贴纸包...请耐心等待", quote=True)
            file = await export_end(uid, target_dir, reply)
            await message.reply_chat_action(ChatAction.UPLOAD_DOCUMENT)
            await message.reply_document(file.as_posix())
        except FileNotFoundError:
            await message.reply("没有选择贴纸，导出失败", quote=True)
        finally:
            await cache.delete(f"sticker:export:{uid}")
            if file:
                file.unlink(missing_ok=True)