From 0ce2893911c20aa435b64bc77359977ecef13363 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 00:25:33 +0800 Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E5=8C=BF=E5=90=8Dcookie=E5=92=8C?= =?UTF-8?q?=E7=94=A8=E6=88=B7cookie=E4=B8=80=E8=B5=B7=E8=B0=83=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 31 +++++++++++++++++++++++- nonebot_bison/config/db_model.py | 5 ++++ nonebot_bison/scheduler/manager.py | 2 ++ nonebot_bison/utils/site.py | 38 +++++++++++++++++++----------- tests/config/test_cookie.py | 4 ++-- 5 files changed, 63 insertions(+), 17 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 909d6f9..fc5ea53 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -272,7 +272,20 @@ class DBConfig: res = [cookie for cookie in res if cookie.id in ids] return res - async def add_cookie(self, platform_name: str, content: str) -> int: + async def get_unviersal_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]: + async with create_session() as sess: + query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 + if platform_name: + query = query.where(Cookie.platform_name == platform_name) + query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) + res = (await sess.scalars(query)).all() + if target: + query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target) + ids = set((await sess.scalars(query)).all()) + res = [cookie for cookie in res if cookie.id in ids] + return res + + async def add_cookie_with_content(self, platform_name: str, content: str) -> int: async with create_session() as sess: cookie = Cookie(platform_name=platform_name, content=content) sess.add(cookie) @@ -280,6 +293,13 @@ class DBConfig: await sess.refresh(cookie) return cookie.id + async def add_cookie(self, cookie: Cookie) -> int: + async with create_session() as sess: + sess.add(cookie) + await sess.commit() + await sess.refresh(cookie) + return cookie.id + async def update_cookie(self, cookie: Cookie): async with create_session() as sess: cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id)) @@ -306,6 +326,15 @@ class DBConfig: ) return list((await sess.scalars(query)).all()) + async def get_universal_cookie(self, platform_name: str) -> list[Cookie]: + async with create_session() as sess: + query = ( + select(Cookie) + .where(Cookie.platform_name == platform_name) + .where(Cookie.is_universal == True) # noqa: E712 + ) + return list((await sess.scalars(query)).all()) + async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): async with create_session() as sess: target_obj = await sess.scalar( diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 4362ee4..590f978 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -74,10 +74,15 @@ class Cookie(Model): id: Mapped[int] = mapped_column(primary_key=True) platform_name: Mapped[str] = mapped_column(String(20)) content: Mapped[str] = mapped_column(String(1024)) + # 最后使用的时刻 last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1)) + # Cookie 当前的状态 status: Mapped[str] = mapped_column(String(20), default="") + # 使用一次之后，需要的冷却时间 cd: Mapped[int] = mapped_column(default=0) + # 是否是通用 Cookie，默认用于匿名 Cookie is_universal: Mapped[bool] = mapped_column(default=False) + # 标签，扩展用 tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={}) targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie") diff --git a/nonebot_bison/scheduler/manager.py b/nonebot_bison/scheduler/manager.py index badb91a..06486eb 100644 --- a/nonebot_bison/scheduler/manager.py +++ b/nonebot_bison/scheduler/manager.py @@ -30,6 +30,8 @@ async def init_scheduler(): else: _schedule_class_platform_dict[site].append(platform_name) for site, target_list in _schedule_class_dict.items(): + if hasattr(site.client_mgr, "_cookie_client_manger_"): + await site.client_mgr.init_universal_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/utils/site.py b/nonebot_bison/utils/site.py index ac3ec94..509fee8 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -42,6 +42,7 @@ class DefaultClientManager(ClientManager): class CookieClientManager(ClientManager): + _cookie_client_manger_ = True _platform_name: str _cookie_cd: int = 10 @@ -49,9 +50,6 @@ class CookieClientManager(ClientManager): """hook函数生成器，用于回写请求状态到数据库""" async def _response_hook(resp: httpx.Response): - if cookie.content == "{}": - logger.debug(f"未携带bison cookie: {resp.request.url}") - return if resp.status_code == 200: logger.debug(f"请求成功: {cookie.id} {resp.request.url}") cookie.status = "success" @@ -63,18 +61,30 @@ class CookieClientManager(ClientManager): return _response_hook + @classmethod + async def init_universal_cookie(cls): + """移除已有的匿名cookie，添加一个新的匿名cookie""" + universal_cookies = await config.get_unviersal_cookie(cls._platform_name) + for cookie in universal_cookies: + await config.delete_cookie(cookie.id) + universal_cookie = Cookie(platform_name=cls._platform_name, content="{}", is_universal=True) + await config.add_cookie(universal_cookie) + + @classmethod + async def init_cookie(cls, cookie: Cookie): + """初始化cookie，添加用户cookie时使用""" + cookie.cd = cls._cookie_cd + async def _choose_cookie(self, target: Target) -> Cookie: """选择 cookie 的具体算法""" if not target: return Cookie(content="{}") - cookies = await config.get_cookie_by_target(target, self._platform_name) - cookies = ( - cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=self._cookie_cd) < datetime.now() - ) - cookies = list(cookies) - if not cookies: - return Cookie(content="{}") - cookie = max(cookies, key=lambda x: x.last_usage) + cookies = [ + *(await config.get_cookie_by_target(target, self._platform_name)), + *(await config.get_universal_cookie(self._platform_name)), + ] + cookies = (cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=cookie.cd) < datetime.now()) + cookie = min(cookies, key=lambda x: x.last_usage) return cookie async def _check_cookie(self, cookie: Cookie) -> Cookie: @@ -85,10 +95,10 @@ class CookieClientManager(ClientManager): """获取 client，根据 target 选择 cookie""" client = http_client() cookie = await self._choose_cookie(target) - if cookie.content != "{}": - logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") + if cookie.is_universal: + logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie") else: - logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie") + logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") return await self._assemble_client(client, cookie) diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py index 92c80d7..683b104 100644 --- a/tests/config/test_cookie.py +++ b/tests/config/test_cookie.py @@ -50,7 +50,7 @@ async def test_cookie_by_user(app: App, init_scheduler): tags=[], ) - await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") + await config.add_cookie_with_content(TargetQQGroup(group_id=123), "weibo", "cookie") cookies = await config.get_cookie(TargetQQGroup(group_id=123)) cookie = cookies[0] @@ -93,7 +93,7 @@ async def test_cookie_target_by_target(app: App, init_scheduler): tags=[], ) - id = await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") + id = await config.add_cookie_with_content(TargetQQGroup(group_id=123), "weibo", "cookie") await config.add_cookie_target(T_Target("weibo_id"), "weibo", id)