From efcf42129010808a1f07ccb6379e120d604d855b Mon Sep 17 00:00:00 2001
From: suyiiyii <suyiiyii@gmail.com>
Date: Wed, 11 Dec 2024 21:10:26 +0800
Subject: [PATCH] =?UTF-8?q?:sparkles:=20B=E7=AB=99=20cookie=E9=80=89?=
 =?UTF-8?q?=E6=8B=A9=E7=AD=96=E7=95=A5=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 nonebot_bison/platform/bilibili/scheduler.py | 39 +++++++++++++++++++-
 nonebot_bison/utils/site.py                  |  4 +-
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/nonebot_bison/platform/bilibili/scheduler.py b/nonebot_bison/platform/bilibili/scheduler.py
index 1dddbec..e7481ba 100644
--- a/nonebot_bison/platform/bilibili/scheduler.py
+++ b/nonebot_bison/platform/bilibili/scheduler.py
@@ -1,15 +1,18 @@
 import json
 import random
+from collections.abc import Callable
 from typing_extensions import override
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, TypeVar
 
-from httpx import AsyncClient
 from nonebot import logger, require
+from httpx import Response, AsyncClient
 from playwright.async_api import Cookie
 
 from nonebot_bison.utils import Site, http_client
 
+from ...config import config
+from ...config.db_model import Target
 from ...utils.site import CookieClientManager
 from ...config.db_model import Cookie as CookieModel
 
@@ -25,6 +28,7 @@ B = TypeVar("B", bound="Bilibili")
 class BilibiliClientManager(CookieClientManager):
 
     _default_cookie_cd = timedelta(seconds=120)
+    _current_user_cookie: CookieModel | None = None
 
     async def _get_cookies(self) -> list[Cookie]:
         browser = await get_browser()
@@ -59,6 +63,39 @@ class BilibiliClientManager(CookieClientManager):
         )
         return cookie
 
+    def _generate_hook(self, cookie: Cookie) -> Callable:
+        """hook 函数生成器，用于回写请求状态到数据库"""
+
+        async def _response_hook(resp: Response):
+            if resp.status_code == 200:
+                logger.trace(f"请求成功: {cookie.id} {resp.request.url}")
+                cookie.status = "success"
+            else:
+                logger.warning(f"请求失败: {cookie.id} {resp.request.url}, 状态码: {resp.status_code}")
+                cookie.status = "failed"
+                self._current_user_cookie = None
+            cookie.last_usage = datetime.now()
+            await config.update_cookie(cookie)
+
+        return _response_hook
+
+    async def _get_next_user_cookie(self) -> CookieModel | None:
+        """选择下一个用户 cookie"""
+        cookies = await config.get_cookie(self._site_name, is_anonymous=False)
+        available_cookies = [cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now()]
+        if not available_cookies:
+            return None
+        cookie = min(available_cookies, key=lambda x: x.last_usage)
+        return cookie
+
+    async def _choose_cookie(self, target: Target | None) -> Cookie:
+        """选择 cookie 的具体算法"""
+        if self._current_user_cookie is None:
+            self._current_user_cookie = await self._get_next_user_cookie()
+        if self._current_user_cookie:
+            return self._current_user_cookie
+        return await config.get_cookie(self._site_name, is_anonymous=True)
+
     @override
     async def refresh_client(self):
         await self._refresh_anonymous_cookie()
diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py
index aecb00f..0f8f2f2 100644
--- a/nonebot_bison/utils/site.py
+++ b/nonebot_bison/utils/site.py
@@ -123,8 +123,8 @@ class CookieClientManager(ClientManager):
     async def _choose_cookie(self, target: Target | None) -> Cookie:
         """选择 cookie 的具体算法"""
         cookies = await config.get_cookie(self._site_name, target)
-        avaliable_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
-        cookie = min(avaliable_cookies, key=lambda x: x.last_usage)
+        available_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
+        cookie = min(available_cookies, key=lambda x: x.last_usage)
         return cookie
 
     async def get_client(self, target: Target | None) -> AsyncClient: