From ca9132a54b6ff65a9d9e78862a9e7db333729336 Mon Sep 17 00:00:00 2001 From: omg-xtao <100690902+omg-xtao@users.noreply.github.com> Date: Mon, 29 Apr 2024 16:52:41 +0800 Subject: [PATCH] :sparkles: Support admin redeem code for users --- plugins/starrail/redeem/redeem.py | 97 +++++++++++++++++++++++++++++-- plugins/starrail/redeem/runner.py | 13 +++-- 2 files changed, 101 insertions(+), 9 deletions(-) diff --git a/plugins/starrail/redeem/redeem.py b/plugins/starrail/redeem/redeem.py index 2279a91..6f511d0 100644 --- a/plugins/starrail/redeem/redeem.py +++ b/plugins/starrail/redeem/redeem.py @@ -1,13 +1,19 @@ import asyncio +import contextlib import time -from typing import List +from asyncio import sleep +from typing import List, Tuple -from telegram import Update -from telegram.error import BadRequest +from simnet import Region +from telegram import Update, Message +from telegram.error import BadRequest, Forbidden from telegram.ext import CallbackContext from telegram.ext import filters from core.plugin import Plugin, handler +from gram_core.basemodel import RegionEnum +from gram_core.services.cookies import CookiesService +from gram_core.services.cookies.models import CookiesStatusEnum from gram_core.services.users.services import UserAdminService from plugins.starrail.redeem.runner import RedeemRunner, RedeemResult, RedeemQueueFull from plugins.tools.genshin import GenshinHelper @@ -19,6 +25,16 @@ REDEEM_TEXT = """#### 兑换结果 #### UID: {} 兑换码：{} 兑换结果：{}""" +REDEEM_ALL_TEXT = """#### 批量兑换 #### + +兑换码：{} +正在兑换中，请稍等 +{} / {}""" +REDEEM_ALL_FAIL_TEXT = """#### 批量兑换 #### + +兑换码：{} +兑换成功：{} +兑换失败：{}""" class Redeem(Plugin): @@ -28,12 +44,14 @@ class Redeem(Plugin): self, genshin_helper: GenshinHelper, user_admin_service: UserAdminService, + cookies_service: CookiesService, ): self.genshin_helper = genshin_helper self.user_admin_service = user_admin_service self.max_code_in_pri_message = 5 self.max_code_in_pub_message = 3 self.redeem_runner = RedeemRunner(genshin_helper) + self.cookies_service = cookies_service async def _callback(self, data: "RedeemResult") -> None: code = data.code @@ -53,13 +71,17 @@ class Redeem(Plugin): if reply_message and filters.ChatType.GROUPS.filter(reply_message): self.add_delete_message_job(reply_message) - async def redeem_one_code(self, update: Update, user_id: int, uid: int, code: str): + async def redeem_one_code(self, update: Update, user_id: int, uid: int, code: str, chinese: bool): if not code: return message = update.effective_message reply_message = await message.reply_text("正在兑换中，请稍等") task_data = RedeemResult(user_id=user_id, code=code, uid=uid, message=reply_message) + if chinese: + task_data.error = "此服务器暂不支持进行兑换哦~" + await self._callback(task_data) + return priority = 1 if await self.user_admin_service.is_admin(user_id) else 2 try: await self.redeem_runner.run(task_data, self._callback, priority) @@ -70,10 +92,11 @@ class Redeem(Plugin): async def redeem_codes(self, update: Update, user_id: int, codes: List[str]): async with self.genshin_helper.genshin(user_id) as client: + chinese = client.region == Region.CHINESE uid = client.player_id tasks = [] for code in codes: - tasks.append(self.redeem_one_code(update, user_id, uid, code)) + tasks.append(self.redeem_one_code(update, user_id, uid, code, chinese)) await asyncio.gather(*tasks) @handler.command(command="redeem", cookie=True, block=False) @@ -98,3 +121,67 @@ class Redeem(Plugin): codes = [i for i in args[0].split("_")[1:] if i][: self.max_code_in_pri_message] logger.info("用户 %s[%s] 通过start命令 进入兑换码兑换流程 codes[%s]", user.full_name, user.id, codes) await self.redeem_codes(update, user.id, codes) + + async def _job_callback(self, data: "RedeemResult") -> None: + if data.error: + logger.warning("执行自动兑换兑换码时发生错误 user_id[%s] message[%s]", data.user_id, data.error) + data.count[1] += 1 + return + data.count[0] += 1 + user_id = data.user_id + code = data.code + uid = data.uid if data.uid else "未知" + msg = "成功" + today = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + + text = REDEEM_TEXT.format(today, uid, code, msg) + try: + await self.application.bot.send_message(user_id, text) + except BadRequest as exc: + logger.warning("执行自动兑换兑换码时发生错误 user_id[%s] Message[%s]", user_id, exc.message) + except Forbidden as exc: + logger.warning("执行自动兑换兑换码时发生错误 user_id[%s] message[%s]", user_id, exc.message) + except Exception as exc: + logger.warning("执行自动兑换兑换码时发生错误 user_id[%s]", user_id, exc_info=exc) + + async def job_redeem_one_code(self, user_id: int, code: str, count: List[int]): + task_data = RedeemResult(user_id=user_id, code=code, count=count) + priority = 1 if await self.user_admin_service.is_admin(user_id) else 2 + try: + await self.redeem_runner.run(task_data, self._job_callback, priority, True) + except RedeemQueueFull: + await sleep(5) + await self.job_redeem_one_code(user_id, code, count) + + async def do_redeem_job(self, message: "Message", code: str) -> Tuple[int, int]: + count = [0, 0] + task_list = await self.cookies_service.get_all( + region=RegionEnum.HOYOLAB, status=CookiesStatusEnum.STATUS_SUCCESS + ) + task_len = len(task_list) + for idx, task_db in enumerate(task_list): + user_id = task_db.user_id + try: + await self.job_redeem_one_code(user_id, code, count) + except Exception as exc: + logger.warning("执行自动兑换兑换码时发生错误 user_id[%s]", user_id, exc_info=exc) + if idx % 10 == 0: + text = REDEEM_ALL_TEXT.format(code, idx, task_len) + with contextlib.suppress(Exception): + await message.edit_text(text) + return count[0], count[1] + + @handler.command(command="redeem_all", admin=True, block=False) + async def redeem_all_command_start(self, update: Update, context: CallbackContext) -> None: + message = update.effective_message + codes = [i for i in self.get_args(context) if i] + self.log_user(update, logger.info, "兑换码批量兑换命令请求 codes[%s]", codes) + if not codes: + await message.reply_text("请输入兑换码") + return + code = codes[0] + reply = await message.reply_text("开始运行批量兑换任务，请等待...") + success, failed = await self.do_redeem_job(reply, code) + text = REDEEM_ALL_FAIL_TEXT.format(code, success, failed) + await message.reply_text(text) + self.add_delete_message_job(reply, delay=1) diff --git a/plugins/starrail/redeem/runner.py b/plugins/starrail/redeem/runner.py index c6b5e66..8d13fad 100644 --- a/plugins/starrail/redeem/runner.py +++ b/plugins/starrail/redeem/runner.py @@ -6,6 +6,7 @@ from typing import Coroutine, Any, Optional, List, TYPE_CHECKING, Union from simnet.errors import RegionNotSupported, RedemptionInvalid, RedemptionClaimed, RedemptionCooldown from telegram import Message +from gram_core.basemodel import RegionEnum from plugins.tools.genshin import GenshinHelper if TYPE_CHECKING: @@ -16,9 +17,10 @@ if TYPE_CHECKING: class RedeemResult: user_id: int code: str - message: Message + message: Optional[Message] = None error: Optional[str] = None uid: Optional[int] = 0 + count: Optional[List[int]] = None class RedeemRunnerTask: @@ -57,8 +59,9 @@ class RedeemRunner: data: RedeemResult, callback_task: "(result: RedeemResult) -> Coroutine[Any, Any, None]", priority: int = 2, + only_region: bool = False, ) -> None: - redeem_task = self.redeem_code(data) + redeem_task = self.redeem_code(data, only_region) queue_task = RedeemRunnerTask(self._execute_queue(redeem_task, callback_task)) if priority == 2 and self.queue.qsize() >= (self.queue_size - 1): raise RedeemQueueFull() @@ -71,10 +74,12 @@ class RedeemRunner: await task.run() await asyncio.sleep(5) - async def redeem_code(self, result: RedeemResult) -> RedeemResult: + async def redeem_code(self, result: RedeemResult, only_region: bool) -> RedeemResult: error = None try: - async with self.genshin_helper.genshin(result.user_id) as client: + async with self.genshin_helper.genshin( + result.user_id, region=RegionEnum.HOYOLAB if only_region else None + ) as client: client: "StarRailClient" result.uid = client.player_id await client.redeem_code_by_hoyolab(result.code)