PamGram/core/services/cookies/services.py

from typing import List, Optional

import genshin
from genshin import Game, GenshinException, InvalidCookies, TooManyRequests, types

from core.base_service import BaseService
from core.basemodel import RegionEnum
from core.services.cookies.cache import PublicCookiesCache
from core.services.cookies.error import CookieServiceError, TooManyRequestPublicCookies
from core.services.cookies.models import CookiesDataBase as Cookies, CookiesStatusEnum
from core.services.cookies.repositories import CookiesRepository
from utils.log import logger

__all__ = ("CookiesService", "PublicCookiesService")


class CookiesService(BaseService):
    def __init__(self, cookies_repository: CookiesRepository) -> None:
        self._repository: CookiesRepository = cookies_repository

    async def update(self, cookies: Cookies):
        await self._repository.update(cookies)

    async def add(self, cookies: Cookies):
        await self._repository.add(cookies)

    async def get(
        self,
        user_id: int,
        account_id: Optional[int] = None,
        region: Optional[RegionEnum] = None,
    ) -> Optional[Cookies]:
        return await self._repository.get(user_id, account_id, region)

    async def delete(self, cookies: Cookies) -> None:
        return await self._repository.delete(cookies)


class PublicCookiesService(BaseService):
    def __init__(self, cookies_repository: CookiesRepository, public_cookies_cache: PublicCookiesCache):
        self._cache = public_cookies_cache
        self._repository: CookiesRepository = cookies_repository
        self.count: int = 0
        self.user_times_limiter = 3 * 3

    async def initialize(self) -> None:
        logger.info("正在初始化公共Cookies池")
        await self.refresh()
        logger.success("刷新公共Cookies池成功")

    async def refresh(self):
        """刷新公共Cookies 定时任务
        :return:
        """
        user_list: List[int] = []
        cookies_list = await self._repository.get_all_by_region(RegionEnum.HYPERION)  # 从数据库获取2
        for cookies in cookies_list:
            if cookies.status is None or cookies.status == CookiesStatusEnum.STATUS_SUCCESS:
                user_list.append(cookies.user_id)
        if len(user_list) > 0:
            add, count = await self._cache.add_public_cookies(user_list, RegionEnum.HYPERION)
            logger.info("国服公共Cookies池已经添加[%s]个 当前成员数为[%s]", add, count)
        user_list.clear()
        cookies_list = await self._repository.get_all_by_region(RegionEnum.HOYOLAB)
        for cookies in cookies_list:
            if cookies.status is None or cookies.status == CookiesStatusEnum.STATUS_SUCCESS:
                user_list.append(cookies.user_id)
        if len(user_list) > 0:
            add, count = await self._cache.add_public_cookies(user_list, RegionEnum.HOYOLAB)
            logger.info("国际服公共Cookies池已经添加[%s]个 当前成员数为[%s]", add, count)

    async def get_cookies(self, user_id: int, region: RegionEnum = RegionEnum.NULL):
        """获取公共Cookies
        :param user_id: 用户ID
        :param region: 注册的服务器
        :return:
        """
        user_times = await self._cache.incr_by_user_times(user_id)
        if int(user_times) > self.user_times_limiter:
            logger.warning("用户 %s 使用公共Cookie次数已经到达上限", user_id)
            raise TooManyRequestPublicCookies(user_id)
        while True:
            public_id, count = await self._cache.get_public_cookies(region)
            cookies = await self._repository.get(public_id, region=region)
            if cookies is None:
                await self._cache.delete_public_cookies(public_id, region)
                continue
            if region == RegionEnum.HYPERION:
                client = genshin.Client(cookies=cookies.data, game=types.Game.GENSHIN, region=types.Region.CHINESE)
            elif region == RegionEnum.HOYOLAB:
                client = genshin.Client(
                    cookies=cookies.data, game=types.Game.GENSHIN, region=types.Region.OVERSEAS, lang="zh-cn"
                )
            else:
                raise CookieServiceError
            try:
                if client.cookie_manager.user_id is None:
                    raise RuntimeError("account_id not found")
                record_card = (await client.get_record_cards())[0]
                if record_card.game == Game.GENSHIN and region == RegionEnum.HYPERION:
                    await client.get_partial_genshin_user(record_card.uid)
            except InvalidCookies as exc:
                if exc.retcode in (10001, -100):
                    logger.warning("用户 [%s] Cookies无效", public_id)
                elif exc.retcode == 10103:
                    logger.warning("用户 [%s] Cookie有效，但没有绑定到游戏帐户", public_id)
                else:
                    logger.warning("Cookies无效 ")
                    logger.exception(exc)
                cookies.status = CookiesStatusEnum.INVALID_COOKIES
                await self._repository.update(cookies)
                await self._cache.delete_public_cookies(cookies.user_id, region)
                continue
            except TooManyRequests:
                logger.warning("用户 [%s] 查询次数太多或操作频繁", public_id)
                cookies.status = CookiesStatusEnum.TOO_MANY_REQUESTS
                await self._repository.update(cookies)
                await self._cache.delete_public_cookies(cookies.user_id, region)
                continue
            except GenshinException as exc:
                if exc.retcode == 1034:
                    logger.warning("用户 [%s] 触发验证", public_id)
                else:
                    logger.warning("用户 [%s] 获取账号信息发生错误，错误信息为", public_id)
                    logger.exception(exc)
                await self._cache.delete_public_cookies(cookies.user_id, region)
                continue
            except RuntimeError as exc:
                if "account_id not found" in str(exc):
                    cookies.status = CookiesStatusEnum.INVALID_COOKIES
                    await self._repository.update(cookies)
                    await self._cache.delete_public_cookies(cookies.user_id, region)
                    continue
                raise exc
            except Exception as exc:
                await self._cache.delete_public_cookies(cookies.user_id, region)
                raise exc
            logger.info("用户 user_id[%s] 请求用户 user_id[%s] 的公共Cookies 该Cookie使用次数为%s次 ", user_id, public_id, count)
            return cookies