From 16331b50d5195cb98c7202a9ced0b19d2aaf335b Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 00:59:19 +0800 Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E6=A0=B9=E6=8D=AE=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=BA=93=E7=9A=84=E4=BF=AE=E6=94=B9=E5=AF=B9=E5=BA=94?= =?UTF-8?q?=E7=9A=84=E9=80=BB=E8=BE=91=EF=BC=8C=E5=90=8C=E6=97=B6=E7=A7=BB?= =?UTF-8?q?=E9=99=A4=E5=A4=9A=E4=BD=99=E7=9A=84init=5Fcookie=E5=92=8C=5Fch?= =?UTF-8?q?eck=5Fcookie?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/scheduler/manager.py | 7 ++++-- nonebot_bison/sub_manager/add_cookie.py | 6 +---- nonebot_bison/utils/site.py | 33 +++++++++++-------------- 3 files changed, 21 insertions(+), 25 deletions(-) diff --git a/nonebot_bison/scheduler/manager.py b/nonebot_bison/scheduler/manager.py index 852f1bb..27a294d 100644 --- a/nonebot_bison/scheduler/manager.py +++ b/nonebot_bison/scheduler/manager.py @@ -1,3 +1,5 @@ +from typing import cast + from nonebot.log import logger from ..utils import Site @@ -7,7 +9,7 @@ from ..config.db_model import Target from ..types import Target as T_Target from ..platform import platform_manager from ..plugin_config import plugin_config -from ..utils.site import is_cookie_client_manager +from ..utils.site import CookieClientManager, is_cookie_client_manager scheduler_dict: dict[type[Site], Scheduler] = {} @@ -32,7 +34,8 @@ async def init_scheduler(): _schedule_class_platform_dict[site].append(platform_name) for site, target_list in _schedule_class_dict.items(): if is_cookie_client_manager(site.client_mgr): - await site.client_mgr.refresh_universal_cookie() + client_mgr = cast(CookieClientManager, site.client_mgr) + await client_mgr.refresh_anonymous_cookie() if not plugin_config.bison_use_browser and site.require_browser: logger.warning(f"{site.name} requires browser, it will not schedule.") continue diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index b7d35ac..20aa004 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -5,8 +5,6 @@ from nonebot.matcher import Matcher from nonebot.params import Arg, ArgPlainText from nonebot.adapters import Message, MessageTemplate -from ..config import config -from ..config.db_model import Cookie from ..platform import platform_manager from .utils import common_platform, gen_handle_cancel from ..utils.site import CookieClientManager, is_cookie_client_manager @@ -64,10 +62,8 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.handle() async def add_cookie_process(state: T_State): - cookie = Cookie(site_name=state["site"].name, content=state["cookie"]) client_mgr = cast(CookieClientManager, platform_manager[state["platform"]].site.client_mgr) - cookie = await client_mgr.init_cookie(cookie) - await config.add_cookie(cookie) + await client_mgr.add_user_cookie(state["cookie"]) await add_cookie.finish( f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" ) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 5c1bb8d..e39d1e9 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -44,26 +44,27 @@ class DefaultClientManager(ClientManager): class CookieClientManager(ClientManager): _site_name: str - _cookie_cd: int = 10 + _default_cd: int = timedelta(seconds=10) @classmethod - async def refresh_universal_cookie(cls): + async def refresh_anonymous_cookie(cls): """移除已有的匿名cookie，添加一个新的匿名cookie""" - universal_cookies = await config.get_unviersal_cookie(cls._site_name) - universal_cookie = Cookie(site_name=cls._site_name, content="{}", is_universal=True, tags={"temporary": True}) - for cookie in universal_cookies: - if not cookie.tags.get("temporary"): + anonymous_cookies = await config.get_unviersal_cookie(cls._site_name) + anonymous_cookie = Cookie(site_name=cls._site_name, content="{}", is_universal=True, is_anonymous=True) + for cookie in anonymous_cookies: + if not cookie.is_anonymous: continue await config.delete_cookie_by_id(cookie.id) - universal_cookie.id = cookie.id # 保持原有的id - await config.add_cookie(universal_cookie) + anonymous_cookie.id = cookie.id # 保持原有的id + anonymous_cookie.last_usage = datetime.now() # 使得第一次请求优先使用用户 cookie + await config.add_cookie(anonymous_cookie) @classmethod - async def init_cookie(cls, cookie: Cookie) -> Cookie: - """初始化 cookie，添加用户 cookie 时使用""" - cookie.cd = cls._cookie_cd - cookie.last_usage = datetime.now() # 使得优先使用用户 cookie - return cookie + async def add_user_cookie(cls, content: str): + """添加用户 cookie""" + cookie = Cookie(site_name=cls._site_name, content=content) + cookie.cd = cls._default_cd + config.add_cookie(cookie) @classmethod async def validate_cookie(cls, content: str) -> bool: @@ -103,14 +104,10 @@ class CookieClientManager(ClientManager): cookies = await config.get_universal_cookie(self._site_name) if target: cookies += await config.get_cookie(self._site_name, target) - cookies = (cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=cookie.cd) < datetime.now()) + cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now()) cookie = min(cookies, key=lambda x: x.last_usage) return cookie - async def _check_cookie(self, cookie: Cookie) -> Cookie: - """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控""" - return cookie - async def get_client(self, target: Target | None) -> AsyncClient: """获取 client，根据 target 选择 cookie""" client = http_client()