✨ Support hoyolab accompany roles

plugins/admin/set_command.py

@@ -33,6 +33,8 @@ class SetCommandPlugin(Plugin):
             BotCommand("daily_note_tasks", "自动便笺提醒"),
             BotCommand("cookies_import", "从其他 BOT 导入账号信息"),
             BotCommand("cookies_export", "导出账号信息给其他 BOT"),
+            BotCommand("redeem", "（国际服）兑换 Key"),
+            BotCommand("accompany", "（国际服）角色陪伴"),
             BotCommand("privacy", "隐私政策"),
         ]
         group_command = [
@@ -76,6 +78,7 @@ class SetCommandPlugin(Plugin):
             BotCommand("add_block", "添加黑名单"),
             BotCommand("del_block", "移除黑名单"),
             BotCommand("signal_log_rank_recount", "重新统计抽卡排行榜"),
+            BotCommand("accompany_all", "全部账号重新角色陪伴"),
         ]
         await context.bot.set_my_commands(
             commands=group_command

plugins/jobs/accompany.py

@@ -0,0 +1,166 @@
+import asyncio
+import datetime
+import random
+import time
+from typing import TYPE_CHECKING, Optional, List
+from contextlib import asynccontextmanager
+
+from simnet import Region
+from simnet.client.components.lab import LabClient
+from simnet.errors import BadRequest as SimnetBadRequest, TimedOut as SimnetTimedOut, InvalidCookies
+from telegram.constants import ParseMode
+from telegram.error import Forbidden, BadRequest
+
+from gram_core.basemodel import RegionEnum
+from gram_core.plugin import Plugin, job, handler
+from gram_core.services.cookies import CookiesService
+from utils.log import logger
+
+if TYPE_CHECKING:
+    from gram_core.services.cookies.models import CookiesDataBase
+    from telegram import Update
+    from telegram.ext import ContextTypes
+    from simnet.models.lab.accompany import AccompanyRole
+
+
+class AccompanySystemError(Exception):
+    def __init__(self, msg: str):
+        self.msg = msg
+
+
+class AccompanySystem(Plugin):
+    def __init__(
+        self,
+        cookies_service: CookiesService,
+    ):
+        self.cookies_service = cookies_service
+        self.accompany_roles = []
+
+    @asynccontextmanager
+    async def client(self, ck: "CookiesDataBase") -> LabClient:
+        stoken = ck.data.get("stoken")
+        if not stoken:
+            raise AccompanySystemError("stoken is None")
+
+        if ck.region == RegionEnum.HYPERION:  # 国服
+            region = Region.CHINESE
+        elif ck.region == RegionEnum.HOYOLAB:  # 国际服
+            region = Region.OVERSEAS
+        else:
+            raise AccompanySystemError("Region is not None")
+
+        async with LabClient(
+            ck.data,
+            region=region,
+            account_id=ck.account_id,
+            lang="zh-cn",
+        ) as client:
+            yield client
+
+    @staticmethod
+    async def request_accompany_role(client: "LabClient", role_id: int, topic_id: int) -> Optional["AccompanyRole"]:
+        try:
+            return await client.request_accompany_role(role_id, topic_id)
+        except SimnetTimedOut:
+            logger.warning("Account[%s] 陪伴 %s 失败，API请求超时", client.account_id, role_id)
+            return None
+        except InvalidCookies as exc:
+            raise exc
+        except SimnetBadRequest as error:
+            logger.warning("Account[%s] 陪伴 %s 失败，API返回信息为 %s", client.account_id, role_id, str(error))
+            return None
+
+    async def start_accompany(
+        self,
+        client: "LabClient",
+        is_sleep: bool = False,
+        is_raise: bool = False,
+    ) -> str:
+        if is_sleep:
+            await asyncio.sleep(random.randint(0, 3))  # nosec
+        success, failed, coins = 0, 0, 0
+        try:
+            if not self.accompany_roles:
+                self.accompany_roles = await client.get_accompany_roles()
+            for role in self.accompany_roles:
+                data = await self.request_accompany_role(client, role.role_id, role.topic_id)
+                if data and data.increase_accompany_point:
+                    success += 1
+                    coins += data.increase_accompany_point
+                else:
+                    failed += 1
+        except SimnetTimedOut as error:
+            logger.warning("Account[%s] 获取陪伴信息失败，API请求超时", client.account_id)
+            if is_raise:
+                raise error
+            return "获取陪伴信息失败，API请求超时"
+        except SimnetBadRequest as error:
+            logger.warning("Account[%s] 获取陪伴信息失败，API返回信息为 %s", client.account_id, str(error))
+            if is_raise:
+                raise error
+            return f"获取陪伴信息失败，API返回信息为 {str(error)}"
+        if success == 0 and coins == 0 and is_raise:
+            raise AccompanySystemError("共获得陪伴值 0 ，可能已经陪伴过了")
+        today = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+        message = (
+            f"#### 角色陪伴 ####\n"
+            f"时间：{today} (UTC+8)\n"
+            f"Account: {client.account_id}\n"
+            f"陪伴成功: {success}\n"
+            f"陪伴失败：{failed}\n"
+            f"共获得陪伴值: {coins}"
+        )
+        return message
+
+    async def _do_accompany_job(
+        self, context: "ContextTypes.DEFAULT_TYPE", accompany_list: List["CookiesDataBase"], is_raise: bool = True
+    ) -> None:
+        for accompany_db in accompany_list:
+            user_id = accompany_db.user_id
+            text = None
+            try:
+                async with self.client(accompany_db) as client:
+                    text = await self.start_accompany(client, is_sleep=True, is_raise=is_raise)
+            except (AccompanySystemError, SimnetTimedOut, SimnetBadRequest):
+                continue
+            except Exception as exc:
+                logger.error("执行自动角色陪伴时发生错误 user_id[%s]", user_id, exc_info=exc)
+            if text:
+                try:
+                    await context.bot.send_message(user_id, text, parse_mode=ParseMode.HTML)
+                except (BadRequest, Forbidden):
+                    continue
+                except Exception as exc:
+                    logger.error("执行自动角色陪伴时发生错误 user_id[%s]", user_id, exc_info=exc)
+
+    async def do_accompany_job(self, context: "ContextTypes.DEFAULT_TYPE") -> None:
+        accompany_list = await self.cookies_service.get_all(region=RegionEnum.HOYOLAB)
+        await self._do_accompany_job(context, accompany_list)
+
+    @job.run_daily(time=datetime.time(hour=1, minute=1, second=0), name="AccompanyJob")
+    async def accompany(self, context: "ContextTypes.DEFAULT_TYPE"):
+        logger.info("正在执行自动角色陪伴")
+        await self.do_accompany_job(context)
+        logger.success("执行自动角色陪伴完成")
+
+    @handler.command(command="accompany_all", block=False, admin=True)
+    async def accompany_all(self, update: "Update", context: "ContextTypes.DEFAULT_TYPE"):
+        user = update.effective_user
+        logger.info("用户 %s[%s] accompany_all 命令请求", user.full_name, user.id)
+        message = update.effective_message
+        reply = await message.reply_text("正在全部重新角色陪伴，请稍后...")
+        await self.do_accompany_job(context)
+        await reply.edit_text("全部账号重新角色陪伴完成")
+
+    @handler.command(command="accompany", cookie=True, block=False)
+    async def accompany_someone(self, update: "Update", context: "ContextTypes.DEFAULT_TYPE"):
+        user_id = await self.get_real_user_id(update)
+        logger.info("用户 %s 请求角色陪伴", user_id)
+        message = update.effective_message
+        reply = await message.reply_text("正在角色陪伴，请稍后...")
+        ck = await self.cookies_service.get_all(user_id=user_id, region=RegionEnum.HOYOLAB)
+        if not ck:
+            await reply.edit_text("未绑定国际服账号")
+            return
+        await self._do_accompany_job(context, ck, is_raise=False)
+        await reply.edit_text("角色陪伴完成")