PamGram/plugins/account/cookies_export.py

from http.cookies import SimpleCookie
from typing import List, TYPE_CHECKING
from uuid import uuid4

from pydantic import BaseModel
from telegram import (
    InlineKeyboardButton,
    SwitchInlineQueryChosenChat,
    InlineKeyboardMarkup,
    InlineQueryResultArticle,
    InputTextMessageContent,
    InlineQueryResultsButton,
)
from telegram.ext import filters
from telegram.error import BadRequest
from telegram.helpers import create_deep_linked_url

from gram_core.basemodel import RegionEnum
from gram_core.config import config
from gram_core.dependence.redisdb import RedisDB
from gram_core.plugin import Plugin, handler
from gram_core.services.cookies import CookiesService
from gram_core.services.devices import DevicesService
from utils.log import logger

try:
    import ujson as jsonlib
except ImportError:
    import json as jsonlib

if TYPE_CHECKING:
    from telegram import Update, InlineQuery
    from telegram.ext import ContextTypes


class InlineCookies(BaseModel):
    account_id: int
    region: RegionEnum
    data: str


class CookiesExport(Plugin):
    def __init__(
        self,
        redis: RedisDB,
        cookies_service: CookiesService,
        devices_service: DevicesService,
    ):
        self.qname = "plugin:cookies_export:"
        self.ex = 5 * 60
        self.client = redis.client
        self.cookies_service = cookies_service
        self.devices_service = devices_service

    def get_cache_key(self, user_id: int) -> str:
        return f"{self.qname}{user_id}"

    async def set_cache(self, user_id: int, data: List[InlineCookies]) -> None:
        if not data:
            return
        new_data = jsonlib.dumps([i.dict() for i in data])
        await self.client.set(self.get_cache_key(user_id), new_data, ex=self.ex)

    async def get_cache(self, user_id: int) -> List[InlineCookies]:
        _data = await self.client.get(self.get_cache_key(user_id))
        if _data is None:
            return []
        data_str = _data.decode("utf-8")
        data = jsonlib.loads(data_str)
        return [InlineCookies(**i) for i in data]

    async def get_all_cookies(self, user_id: int) -> List[InlineCookies]:
        cookies = await self.cookies_service.get_all(user_id)
        if not cookies:
            return []
        cookies_list = []
        for cookie in cookies:
            if cookie.region not in [RegionEnum.HYPERION, RegionEnum.HOYOLAB]:
                continue
            cookies = SimpleCookie()
            for key, value in cookie.data.items():
                cookies[key] = value

            device = None
            if cookie.region == RegionEnum.HYPERION:
                device = await self.devices_service.get(cookie.account_id)
            if device is not None:
                cookies["x-rpc-device_id"] = device.device_id
                cookies["x-rpc-device_fp"] = device.device_fp
            cookie_str = cookies.output(header="", sep=";")

            cookies_list.append(
                InlineCookies(
                    account_id=cookie.account_id,
                    region=cookie.region,
                    data=cookie_str,
                )
            )
        await self.set_cache(user_id, cookies_list)
        return cookies_list

    @handler.command("cookies_export", block=False)
    @handler.command("start", filters=filters.Regex("cookies_export$"), block=False)
    async def cookies_export(self, update: "Update", _: "ContextTypes.DEFAULT_TYPE"):
        message = update.effective_message
        user = update.effective_user
        logger.info("用户 %s[%s] cookies_export 命令请求", message.from_user.full_name, message.from_user.id)
        data = await self.get_all_cookies(user.id)
        if not data:
            await message.reply_text("没有查询到任何账号信息")
            return
        text = "请点击下方按钮导出账号信息到指定 BOT"
        buttons = [
            [
                InlineKeyboardButton(
                    "选择需要导入账号的 BOT",
                    switch_inline_query_chosen_chat=SwitchInlineQueryChosenChat(
                        query="cookies_export", allow_bot_chats=True
                    ),
                )
            ]
        ]
        await message.reply_text(text, reply_markup=InlineKeyboardMarkup(buttons))

    def gen_cookies_import_buttons(self):
        official_bots = config.bot.official.copy()
        lower_official_bots = [i.lower() for i in official_bots]
        bot_username_lower = self.application.bot.username.lower()
        if bot_username_lower in lower_official_bots:
            official_bots.pop(lower_official_bots.index(bot_username_lower))
        return [
            [
                InlineKeyboardButton(
                    text=name,
                    url=create_deep_linked_url(name, "cookies_export"),
                )
            ]
            for name in official_bots
        ]

    @handler.command("cookies_import", block=False)
    @handler.command("start", filters=filters.Regex("cookies_import$"), block=False)
    async def cookies_import(self, update: "Update", _: "ContextTypes.DEFAULT_TYPE"):
        message = update.effective_message
        user = update.effective_user
        logger.info("用户 %s[%s] cookies_import 命令请求", user.full_name, user.id)
        text = "请点击下方按钮选择您已经绑定了账号的 BOT"
        buttons = self.gen_cookies_import_buttons()
        if not buttons:
            await message.reply_text("没有可用的BOT")
            return
        await message.reply_text(text, reply_markup=InlineKeyboardMarkup(buttons))

    @handler.inline_query(pattern="^cookies_export$", block=False)
    async def inline_query(self, update: "Update", _: "ContextTypes.DEFAULT_TYPE") -> None:
        user = update.effective_user
        ilq: "InlineQuery" = update.inline_query
        cache_data = await self.get_cache(user.id)
        results_list = []
        if not cache_data:
            results_list.append(
                InlineQueryResultArticle(
                    id=str(uuid4()),
                    title="无法导出 Cookies",
                    description="请先使用命令 /cookies_export",
                    input_message_content=InputTextMessageContent("/cookies_export"),
                )
            )
        else:
            name_map = {RegionEnum.HYPERION: "米游社", RegionEnum.HOYOLAB: "HoYoLab"}
            for cookie in cache_data:
                region = name_map[cookie.region]
                results_list.append(
                    InlineQueryResultArticle(
                        id=str(uuid4()),
                        title=f"{region} - {cookie.account_id}",
                        description=f"导出账号ID {cookie.account_id} 的 Cookies",
                        input_message_content=InputTextMessageContent(f"/setcookies {region} {cookie.data}"),
                    )
                )
        try:
            await ilq.answer(
                results=results_list,
                cache_time=0,
                auto_pagination=True,
                button=InlineQueryResultsButton(
                    text="！！导出到不信任对话将有盗号风险！!",
                    start_parameter="cookies_export",
                ),
            )
        except BadRequest as exc:
            if "Query is too old" in exc.message:  # 过时请求全部忽略
                logger.warning("用户 %s[%s] inline_query 请求过时", user.full_name, user.id)
                return
            if "can't parse entities" not in exc.message:
                raise exc
            logger.warning("inline_query发生BadRequest错误", exc_info=exc)
            await ilq.answer(
                results=[],
                button=InlineQueryResultsButton(
                    text="糟糕，发生错误了。",
                    start_parameter="inline_message",
                ),
            )