From 7913f7485a8ca9d5d15cfeb0cb9f2b55b916dc5c Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 22 Aug 2024 20:55:39 +0800 Subject: [PATCH 01/51] =?UTF-8?q?:sparkles:=20=E6=B7=BB=E5=8A=A0cookie?= =?UTF-8?q?=E7=9B=B8=E5=85=B3=E7=9A=84=E6=95=B0=E6=8D=AE=E5=BA=93=E8=A1=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 71 +++++++++++- nonebot_bison/config/db_model.py | 27 ++++- .../migrations/bdec8925b540_add_cookie.py | 63 +++++++++++ nonebot_bison/sub_manager/add_cookie.py | 0 tests/config/__init__.py | 0 tests/config/test_cookie.py | 107 ++++++++++++++++++ 6 files changed, 266 insertions(+), 2 deletions(-) create mode 100644 nonebot_bison/config/migrations/bdec8925b540_add_cookie.py create mode 100644 nonebot_bison/sub_manager/add_cookie.py create mode 100644 tests/config/__init__.py create mode 100644 tests/config/test_cookie.py diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 157b1ef..444d9d9 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -13,7 +13,7 @@ from nonebot_plugin_datastore import create_session from ..types import Tag from ..types import Target as T_Target from .utils import NoSuchTargetException -from .db_model import User, Target, Subscribe, ScheduleTimeWeight +from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp @@ -259,5 +259,74 @@ class DBConfig: ) return res + async def get_cookie_by_user(self, user: PlatformTarget) -> list[Cookie]: + async with create_session() as sess: + res = await sess.scalar( + select(User) + .where(User.user_target == model_dump(user)) + .join(Cookie) + .outerjoin(CookieTarget) + .options(selectinload(User.cookies)) + ) + if not res: + return [] + return res.cookies + + async def add_cookie(self, user: PlatformTarget, platform_name: str, content: str) -> int: + async with create_session() as sess: + user_obj = await sess.scalar(select(User).where(User.user_target == model_dump(user))) + cookie = Cookie(user=user_obj, platform_name=platform_name, content=content) + 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)) + if not cookie_in_db: + return + cookie_in_db.content = cookie.content + cookie_in_db.last_usage = cookie.last_usage + cookie_in_db.status = cookie.status + cookie_in_db.tags = cookie.tags + await sess.commit() + + async def delete_cookie(self, cookie_id: int): + async with create_session() as sess: + await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) + await sess.commit() + + async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> list[Cookie]: + async with create_session() as sess: + query = ( + select(Cookie) + .join(CookieTarget) + .join(Target) + .where(Target.platform_name == platform_name, Target.target == target) + ) + 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( + select(Target).where(Target.platform_name == platform_name, Target.target == target) + ) + cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) + sess.add(cookie_target) + await sess.commit() + + async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): + async with create_session() as sess: + target_obj = await sess.scalar( + select(Target).where(Target.platform_name == platform_name, Target.target == target) + ) + cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + await sess.execute( + delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) + ) + await sess.commit() + config = DBConfig() diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 849094d..603ae21 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -1,4 +1,5 @@ import datetime +from typing import Any from pathlib import Path from nonebot_plugin_saa import PlatformTarget @@ -6,7 +7,7 @@ from sqlalchemy.dialects.postgresql import JSONB from nonebot.compat import PYDANTIC_V2, ConfigDict from nonebot_plugin_datastore import get_plugin_data from sqlalchemy.orm import Mapped, relationship, mapped_column -from sqlalchemy import JSON, String, ForeignKey, UniqueConstraint +from sqlalchemy import JSON, String, DateTime, ForeignKey, UniqueConstraint from ..types import Tag, Category @@ -19,6 +20,7 @@ class User(Model): user_target: Mapped[dict] = mapped_column(JSON().with_variant(JSONB, "postgresql")) subscribes: Mapped[list["Subscribe"]] = relationship(back_populates="user") + cookies: Mapped[list["Cookie"]] = relationship(back_populates="user") @property def saa_target(self) -> PlatformTarget: @@ -36,6 +38,7 @@ class Target(Model): subscribes: Mapped[list["Subscribe"]] = relationship(back_populates="target") time_weight: Mapped[list["ScheduleTimeWeight"]] = relationship(back_populates="target") + cookies: Mapped[list["CookieTarget"]] = relationship(back_populates="target") class ScheduleTimeWeight(Model): @@ -66,3 +69,25 @@ class Subscribe(Model): target: Mapped[Target] = relationship(back_populates="subscribes") user: Mapped[User] = relationship(back_populates="subscribes") + + +class Cookie(Model): + id: Mapped[int] = mapped_column(primary_key=True) + user_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_user.id")) + 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)) + status: Mapped[str] = mapped_column(String(20), default="") + tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={}) + + user: Mapped[User] = relationship(back_populates="cookies") + targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie") + + +class CookieTarget(Model): + id: Mapped[int] = mapped_column(primary_key=True) + target_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_target.id", ondelete="CASCADE")) + cookie_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_cookie.id", ondelete="CASCADE")) + + target: Mapped[Target] = relationship(back_populates="cookies") + cookie: Mapped[Cookie] = relationship(back_populates="targets") diff --git a/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py b/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py new file mode 100644 index 0000000..ac650f3 --- /dev/null +++ b/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py @@ -0,0 +1,63 @@ +"""empty message + +Revision ID: bdec8925b540 +Revises: f9baef347cc8 +Create Date: 2024-08-22 20:53:08.850051 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import Text +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "bdec8925b540" +down_revision = "f9baef347cc8" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "nonebot_bison_cookie", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("user_id", sa.Integer(), nullable=False), + sa.Column("platform_name", sa.String(length=20), nullable=False), + sa.Column("content", sa.String(length=1024), nullable=False), + sa.Column("last_usage", sa.DateTime(), nullable=False), + sa.Column("status", sa.String(length=20), nullable=False), + sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False), + sa.ForeignKeyConstraint( + ["user_id"], ["nonebot_bison_user.id"], name=op.f("fk_nonebot_bison_cookie_user_id_nonebot_bison_user") + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")), + ) + op.create_table( + "nonebot_bison_cookietarget", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("target_id", sa.Integer(), nullable=False), + sa.Column("cookie_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["cookie_id"], + ["nonebot_bison_cookie.id"], + name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["target_id"], + ["nonebot_bison_target.id"], + name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("nonebot_bison_cookietarget") + op.drop_table("nonebot_bison_cookie") + # ### end Alembic commands ### diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/config/__init__.py b/tests/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py new file mode 100644 index 0000000..4515078 --- /dev/null +++ b/tests/config/test_cookie.py @@ -0,0 +1,107 @@ +import datetime + +from nonebug import App + + +async def test_get_platform_target(app: App, init_scheduler): + from nonebot_plugin_saa import TargetQQGroup + + from nonebot_bison.config.db_config import config + from nonebot_bison.types import Target as T_Target + + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=T_Target("weibo_id"), + target_name="weibo_name", + platform_name="weibo", + cats=[], + tags=[], + ) + # await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") + # cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + # + # res = await config.get_platform_target("weibo") + # assert len(res) == 2 + # await config.del_subscribe(TargetQQGroup(group_id=123), T_Target("weibo_id1"), "weibo") + # res = await config.get_platform_target("weibo") + # assert len(res) == 2 + # await config.del_subscribe(TargetQQGroup(group_id=123), T_Target("weibo_id"), "weibo") + # res = await config.get_platform_target("weibo") + # assert len(res) == 1 + # + # async with AsyncSession(get_engine()) as sess: + # res = await sess.scalars(select(Target).where(Target.platform_name == "weibo")) + # assert len(res.all()) == 2 + # await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + + +async def test_cookie_by_user(app: App, init_scheduler): + from nonebot_plugin_saa import TargetQQGroup + + from nonebot_bison.config.db_config import config + from nonebot_bison.types import Target as T_Target + + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=T_Target("weibo_id"), + target_name="weibo_name", + platform_name="weibo", + cats=[], + tags=[], + ) + + await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") + + cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + cookie = cookies[0] + assert len(cookies) == 1 + assert cookie.content == "cookie" + assert cookie.platform_name == "weibo" + cookie.last_usage = 0 + assert cookie.status == "" + assert cookie.tags == {} + cookie.content = "cookie1" + cookie.last_usage = datetime.datetime(2024, 8, 22, 0, 0, 0) + cookie.status = "status1" + cookie.tags = {"tag1": "value1"} + await config.update_cookie(cookie) + cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + + assert len(cookies) == 1 + assert cookies[0].content == cookie.content + assert cookies[0].last_usage == cookie.last_usage + assert cookies[0].status == cookie.status + assert cookies[0].tags == cookie.tags + + await config.delete_cookie(cookies[0].id) + cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + assert len(cookies) == 0 + + +async def test_cookie_target_by_target(app: App, init_scheduler): + from nonebot_plugin_saa import TargetQQGroup + + from nonebot_bison.config.db_config import config + from nonebot_bison.types import Target as T_Target + + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=T_Target("weibo_id"), + target_name="weibo_name", + platform_name="weibo", + cats=[], + tags=[], + ) + + id = await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") + + await config.add_cookie_target(T_Target("weibo_id"), "weibo", id) + + cookies = await config.get_cookie_by_target(T_Target("weibo_id"), "weibo") + assert len(cookies) == 1 + assert cookies[0].content == "cookie" + assert cookies[0].platform_name == "weibo" + + await config.delete_cookie_target(T_Target("weibo_id"), "weibo", id) + cookies = await config.get_cookie_by_target(T_Target("weibo_id"), "weibo") + assert len(cookies) == 0 From c264ad374b22bd57f32a36934a40d07b5625b135 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 22 Aug 2024 21:50:25 +0800 Subject: [PATCH 02/51] :bug: stash --- nonebot_bison/sub_manager/__init__.py | 12 ++++- nonebot_bison/sub_manager/add_cookie.py | 64 +++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 119a0e9..93afa17 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -14,6 +14,7 @@ from nonebot.adapters.onebot.v11.event import PrivateMessageEvent from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub +from .add_cookie import do_add_cookie from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info add_sub_matcher = on_command( @@ -26,12 +27,10 @@ add_sub_matcher = on_command( add_sub_matcher.handle()(set_target_user_info) do_add_sub(add_sub_matcher) - query_sub_matcher = on_command("查询订阅", rule=configurable_to_me, priority=5, block=True) query_sub_matcher.handle()(set_target_user_info) do_query_sub(query_sub_matcher) - del_sub_matcher = on_command( "删除订阅", rule=configurable_to_me, @@ -42,6 +41,15 @@ del_sub_matcher = on_command( del_sub_matcher.handle()(set_target_user_info) do_del_sub(del_sub_matcher) +add_cookie_matcher = on_command( + "添加cookie", + rule=configurable_to_me, + permission=admin_permission(), + priority=5, + block=True, +) +do_add_cookie(add_cookie_matcher) + group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) group_handle_cancel = gen_handle_cancel(group_manage_matcher, "已取消") diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index e69de29..6f9dfd0 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -0,0 +1,64 @@ +from nonebot.typing import T_State +from nonebot.matcher import Matcher +from nonebot.params import Arg, ArgPlainText +from nonebot_plugin_saa import PlatformTarget +from nonebot.adapters import Message, MessageTemplate + +from ..types import Target +from ..config import config +from ..apis import check_sub_target +from ..platform import platform_manager +from .utils import common_platform, ensure_user_info, gen_handle_cancel + + +def do_add_cookie(add_cookie: type[Matcher]): + handle_cancel = gen_handle_cancel(add_cookie, "已中止添加cookie") + + add_cookie.handle()(ensure_user_info(add_cookie)) + + @add_cookie.handle() + async def init_promote(state: T_State): + state["_prompt"] = ( + "请输入想要添加 Cookie 的平台，目前支持，请输入冒号左边的名称：\n" + + "".join( + [f"{platform_name}: {platform_manager[platform_name].name}\n" for platform_name in common_platform] + ) + + "要查看全部平台请输入：“全部”\n中止添加cookie过程请输入：“取消”" + ) + + @add_cookie.got("platform", MessageTemplate("{_prompt}"), [handle_cancel]) + async def parse_platform(state: T_State, platform: str = ArgPlainText()) -> None: + if platform == "全部": + message = "全部平台\n" + "\n".join( + [f"{platform_name}: {platform.name}" for platform_name, platform in platform_manager.items()] + ) + await add_cookie.reject(message) + elif platform == "取消": + await add_cookie.finish("已中止添加cookie") + elif platform in platform_manager: + state["platform"] = platform + else: + await add_cookie.reject("平台输入错误") + + @add_cookie.handle() + async def prepare_get_id(matcher: Matcher, state: T_State): + cur_platform = platform_manager[state["platform"]] + if cur_platform.has_target: + state["_prompt"] = ( + ("1." + cur_platform.parse_target_promot + "\n2.") if cur_platform.parse_target_promot else "" + ) + "请输入 Cookie \n查询id获取方法请回复:“查询”" + else: + matcher.set_arg("cookie", None) # type: ignore + state["id"] = "default" + state["name"] = await check_sub_target(state["platform"], Target("")) + + @add_cookie.got("cookie", MessageTemplate("{_prompt}"), [handle_cancel]) + async def got_id(state: T_State, cookie: Message = Arg()): + cookie_text = cookie.extract_plain_text() + state["cookie"] = cookie_text + state["name"] = await check_sub_target(state["platform"], Target("")) + + @add_cookie.handle() + async def add_cookie(state: T_State, user: PlatformTarget = Arg("target_user_info")): + await config.add_cookie(user, state["platform"], state["cookie"]) + await add_cookie.finish(f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}") From b655eff7551106d62dbd92cf35e5846c10fcb707 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 26 Aug 2024 10:37:03 +0800 Subject: [PATCH 03/51] =?UTF-8?q?:sparkles:=20=E6=94=AF=E6=8C=81=E5=AF=B9?= =?UTF-8?q?=E8=AF=9D=E6=B7=BB=E5=8A=A0cookie?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/apis.py | 5 +++++ nonebot_bison/sub_manager/__init__.py | 1 + nonebot_bison/sub_manager/add_cookie.py | 13 +++++-------- nonebot_bison/sub_manager/add_cookie_target.py | 0 4 files changed, 11 insertions(+), 8 deletions(-) create mode 100644 nonebot_bison/sub_manager/add_cookie_target.py diff --git a/nonebot_bison/apis.py b/nonebot_bison/apis.py index f13580b..e79474f 100644 --- a/nonebot_bison/apis.py +++ b/nonebot_bison/apis.py @@ -10,3 +10,8 @@ async def check_sub_target(platform_name: str, target: Target): client = await scheduler.client_mgr.get_query_name_client() return await platform_manager[platform_name].get_target_name(client, target) + + +async def check_sub_target_cookie(platform_name: str, target: Target, cookie: str): + # TODO + return "check pass" diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 93afa17..4a05ba4 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -48,6 +48,7 @@ add_cookie_matcher = on_command( priority=5, block=True, ) +add_cookie_matcher.handle()(set_target_user_info) do_add_cookie(add_cookie_matcher) group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 6f9dfd0..294cff8 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -6,8 +6,8 @@ from nonebot.adapters import Message, MessageTemplate from ..types import Target from ..config import config -from ..apis import check_sub_target from ..platform import platform_manager +from ..apis import check_sub_target_cookie from .utils import common_platform, ensure_user_info, gen_handle_cancel @@ -44,21 +44,18 @@ def do_add_cookie(add_cookie: type[Matcher]): async def prepare_get_id(matcher: Matcher, state: T_State): cur_platform = platform_manager[state["platform"]] if cur_platform.has_target: - state["_prompt"] = ( - ("1." + cur_platform.parse_target_promot + "\n2.") if cur_platform.parse_target_promot else "" - ) + "请输入 Cookie \n查询id获取方法请回复:“查询”" + state["_prompt"] = "请输入 Cookie" else: matcher.set_arg("cookie", None) # type: ignore state["id"] = "default" - state["name"] = await check_sub_target(state["platform"], Target("")) @add_cookie.got("cookie", MessageTemplate("{_prompt}"), [handle_cancel]) - async def got_id(state: T_State, cookie: Message = Arg()): + async def got_cookie(state: T_State, cookie: Message = Arg()): cookie_text = cookie.extract_plain_text() state["cookie"] = cookie_text - state["name"] = await check_sub_target(state["platform"], Target("")) + state["name"] = await check_sub_target_cookie(state["platform"], Target(""), cookie_text) @add_cookie.handle() - async def add_cookie(state: T_State, user: PlatformTarget = Arg("target_user_info")): + async def add_cookie_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): await config.add_cookie(user, state["platform"], state["cookie"]) await add_cookie.finish(f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}") diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py new file mode 100644 index 0000000..e69de29 From 6f20dbf3580dc76f92747f9311a28c0678b32b22 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 26 Aug 2024 17:28:59 +0800 Subject: [PATCH 04/51] =?UTF-8?q?:sparkles:=20=E6=94=AF=E6=8C=81=E5=AF=B9?= =?UTF-8?q?=E8=AF=9D=E5=85=B3=E8=81=94cookie=E5=88=B0=E8=AE=A2=E9=98=85?= =?UTF-8?q?=E7=9B=AE=E6=A0=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/apis.py | 6 ++ nonebot_bison/config/db_config.py | 29 ++++++- nonebot_bison/config/utils.py | 4 + nonebot_bison/sub_manager/__init__.py | 11 +++ nonebot_bison/sub_manager/add_cookie.py | 4 +- .../sub_manager/add_cookie_target.py | 87 +++++++++++++++++++ tests/config/test_cookie.py | 6 +- 7 files changed, 141 insertions(+), 6 deletions(-) diff --git a/nonebot_bison/apis.py b/nonebot_bison/apis.py index e79474f..9bb9176 100644 --- a/nonebot_bison/apis.py +++ b/nonebot_bison/apis.py @@ -1,4 +1,5 @@ from .types import Target +from .config.db_model import Cookie from .scheduler import scheduler_dict from .platform import platform_manager @@ -15,3 +16,8 @@ async def check_sub_target(platform_name: str, target: Target): async def check_sub_target_cookie(platform_name: str, target: Target, cookie: str): # TODO return "check pass" + + +async def get_cookie_friendly_name(cookie: Cookie): + # TODO + return cookie.platform_name + cookie.content[:10] diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 444d9d9..97cc0db 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -12,7 +12,7 @@ from nonebot_plugin_datastore import create_session from ..types import Tag from ..types import Target as T_Target -from .utils import NoSuchTargetException +from .utils import NoSuchTargetException, DuplicateCookieTargetException from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp @@ -259,12 +259,31 @@ class DBConfig: ) return res - async def get_cookie_by_user(self, user: PlatformTarget) -> list[Cookie]: + async def get_cookie( + self, user: PlatformTarget = None, platform_name: str = None, target: T_Target = None + ) -> list[Cookie]: + async with create_session() as sess: + query = select(Cookie).distinct().join(User) + if user: + user_id = await sess.scalar(select(User.id).where(User.user_target == model_dump(user))) + query = query.where(Cookie.user_id == user_id) + 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 get_cookie_by_user_and_platform(self, user: PlatformTarget, platform_name: str) -> list[Cookie]: async with create_session() as sess: res = await sess.scalar( select(User) .where(User.user_target == model_dump(user)) .join(Cookie) + .where(Cookie.platform_name == platform_name) .outerjoin(CookieTarget) .options(selectinload(User.cookies)) ) @@ -312,6 +331,12 @@ class DBConfig: target_obj = await sess.scalar( select(Target).where(Target.platform_name == platform_name, Target.target == target) ) + # check if relation exists + cookie_target = await sess.scalar( + select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) + ) + if cookie_target: + raise DuplicateCookieTargetException() cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) sess.add(cookie_target) diff --git a/nonebot_bison/config/utils.py b/nonebot_bison/config/utils.py index 8c06497..d62ad76 100644 --- a/nonebot_bison/config/utils.py +++ b/nonebot_bison/config/utils.py @@ -8,3 +8,7 @@ class NoSuchSubscribeException(Exception): class NoSuchTargetException(Exception): pass + + +class DuplicateCookieTargetException(Exception): + pass diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 4a05ba4..b16dcf7 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -15,6 +15,7 @@ from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub from .add_cookie import do_add_cookie +from .add_cookie_target import do_add_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info add_sub_matcher = on_command( @@ -51,6 +52,16 @@ add_cookie_matcher = on_command( add_cookie_matcher.handle()(set_target_user_info) do_add_cookie(add_cookie_matcher) +add_cookie_target_matcher = on_command( + "关联cookie", + rule=configurable_to_me, + permission=admin_permission(), + priority=5, + block=True, +) +add_cookie_target_matcher.handle()(set_target_user_info) +do_add_cookie_target(add_cookie_target_matcher) + group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) group_handle_cancel = gen_handle_cancel(group_manage_matcher, "已取消") diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 294cff8..1f7a525 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -58,4 +58,6 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.handle() async def add_cookie_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): await config.add_cookie(user, state["platform"], state["cookie"]) - await add_cookie.finish(f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}") + await add_cookie.finish( + f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 cookie 关联订阅" + ) diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index e69de29..24196c6 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -0,0 +1,87 @@ +from nonebot.typing import T_State +from nonebot.matcher import Matcher +from nonebot.params import Arg, ArgPlainText +from nonebot.internal.adapter import MessageTemplate +from nonebot_plugin_saa import MessageFactory, PlatformTarget + +from ..config import config +from ..types import Category +from ..utils import parse_text +from ..platform import platform_manager +from ..apis import get_cookie_friendly_name +from .utils import ensure_user_info, gen_handle_cancel + + +def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): + handle_cancel = gen_handle_cancel(add_cookie_target_matcher, "已中止关联 cookie") + + add_cookie_target_matcher.handle()(ensure_user_info(add_cookie_target_matcher)) + + @add_cookie_target_matcher.handle() + async def init_promote(state: T_State, user_info: PlatformTarget = Arg("target_user_info")): + sub_list = await config.list_subscribe(user_info) + if not sub_list: + await add_cookie_target_matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅") + res = "订阅的帐号为：\n" + state["sub_table"] = {} + for index, sub in enumerate(sub_list, 1): + state["sub_table"][index] = { + "platform_name": sub.target.platform_name, + "target": sub.target.target, + } + res += f"{index} {sub.target.platform_name} {sub.target.target_name} {sub.target.target}\n" + if platform := platform_manager.get(sub.target.platform_name): + if platform.categories: + res += " [{}]".format(", ".join(platform.categories[Category(x)] for x in sub.categories)) + if platform.enable_tag: + res += " {}".format(", ".join(sub.tags)) + else: + res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）" + if res[-1] != "\n": + res += "\n" + res += "请输入要关联 cookie 的订阅的序号\n输入'取消'中止" + await MessageFactory(await parse_text(res)).send() + + @add_cookie_target_matcher.got("target_idx", parameterless=[handle_cancel]) + async def got_target_idx(state: T_State, target_idx: str = ArgPlainText()): + try: + target_idx = int(target_idx) + state["target"] = state["sub_table"][target_idx] + except Exception: + await add_cookie_target_matcher.reject("序号错误") + + @add_cookie_target_matcher.handle() + async def init_promote_cookie(state: T_State): + + cookies = await config.get_cookie( + user=state["target_user_info"], platform_name=state["target"]["platform_name"] + ) + associated_cookies = await config.get_cookie( + user=state["target_user_info"], + target=state["target"]["target"], + platform_name=state["target"]["platform_name"], + ) + associated_cookie_ids = {cookie.id for cookie in associated_cookies} + cookies = [cookie for cookie in cookies if cookie.id not in associated_cookie_ids] + if not cookies: + await add_cookie_target_matcher.finish("当前平台暂无 Cookie，请使用“添加cookie”命令添加") + state["cookies"] = cookies + state["_prompt"] = "请选择一个 Cookie，已关联的 Cookie 不会显示\n" + "\n".join( + [f"{idx}. {await get_cookie_friendly_name(cookie)}" for idx, cookie in enumerate(cookies, 1)] + ) + + @add_cookie_target_matcher.got("cookie_idx", MessageTemplate("{_prompt}"), [handle_cancel]) + async def got_cookie_idx(state: T_State, cookie_idx: str = ArgPlainText()): + try: + cookie_idx = int(cookie_idx) + state["cookie"] = state["cookies"][cookie_idx - 1] + except Exception: + await add_cookie_target_matcher.reject("序号错误") + + @add_cookie_target_matcher.handle() + async def add_cookie_target_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): + await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id) + await add_cookie_target_matcher.finish( + f"已关联 Cookie: {await get_cookie_friendly_name(state['cookie'])} " + f"到订阅 {state['target']['platform_name']} {state['target']['target']}" + ) diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py index 4515078..92c80d7 100644 --- a/tests/config/test_cookie.py +++ b/tests/config/test_cookie.py @@ -52,7 +52,7 @@ async def test_cookie_by_user(app: App, init_scheduler): await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") - cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + cookies = await config.get_cookie(TargetQQGroup(group_id=123)) cookie = cookies[0] assert len(cookies) == 1 assert cookie.content == "cookie" @@ -65,7 +65,7 @@ async def test_cookie_by_user(app: App, init_scheduler): cookie.status = "status1" cookie.tags = {"tag1": "value1"} await config.update_cookie(cookie) - cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + cookies = await config.get_cookie(TargetQQGroup(group_id=123)) assert len(cookies) == 1 assert cookies[0].content == cookie.content @@ -74,7 +74,7 @@ async def test_cookie_by_user(app: App, init_scheduler): assert cookies[0].tags == cookie.tags await config.delete_cookie(cookies[0].id) - cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + cookies = await config.get_cookie(TargetQQGroup(group_id=123)) assert len(cookies) == 0 From ffae6f2ec58d6398ebdb631119b7e645012153f4 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 26 Aug 2024 18:09:26 +0800 Subject: [PATCH 05/51] =?UTF-8?q?:sparkles:=20=E6=B7=BB=E5=8A=A0cookie?= =?UTF-8?q?=E7=9A=84=E6=97=B6=E5=80=99=E6=98=BE=E7=A4=BA=E5=85=B3=E8=81=94?= =?UTF-8?q?=E7=9A=84cookie?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/apis.py | 2 +- nonebot_bison/sub_manager/add_cookie.py | 2 +- .../sub_manager/add_cookie_target.py | 31 ++++---------- nonebot_bison/sub_manager/utils.py | 41 ++++++++++++++++++- 4 files changed, 49 insertions(+), 27 deletions(-) diff --git a/nonebot_bison/apis.py b/nonebot_bison/apis.py index 9bb9176..c4f247a 100644 --- a/nonebot_bison/apis.py +++ b/nonebot_bison/apis.py @@ -20,4 +20,4 @@ async def check_sub_target_cookie(platform_name: str, target: Target, cookie: st async def get_cookie_friendly_name(cookie: Cookie): # TODO - return cookie.platform_name + cookie.content[:10] + return f"{cookie.platform_name} [{cookie.content[:10]}]" diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 1f7a525..f6ba356 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -59,5 +59,5 @@ def do_add_cookie(add_cookie: type[Matcher]): async def add_cookie_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): await config.add_cookie(user, state["platform"], state["cookie"]) await add_cookie.finish( - f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 cookie 关联订阅" + f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" ) diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index 24196c6..7244736 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -5,11 +5,9 @@ from nonebot.internal.adapter import MessageTemplate from nonebot_plugin_saa import MessageFactory, PlatformTarget from ..config import config -from ..types import Category from ..utils import parse_text -from ..platform import platform_manager from ..apis import get_cookie_friendly_name -from .utils import ensure_user_info, gen_handle_cancel +from .utils import ensure_user_info, gen_handle_cancel, generate_sub_list_text def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @@ -19,26 +17,9 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def init_promote(state: T_State, user_info: PlatformTarget = Arg("target_user_info")): - sub_list = await config.list_subscribe(user_info) - if not sub_list: - await add_cookie_target_matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅") - res = "订阅的帐号为：\n" - state["sub_table"] = {} - for index, sub in enumerate(sub_list, 1): - state["sub_table"][index] = { - "platform_name": sub.target.platform_name, - "target": sub.target.target, - } - res += f"{index} {sub.target.platform_name} {sub.target.target_name} {sub.target.target}\n" - if platform := platform_manager.get(sub.target.platform_name): - if platform.categories: - res += " [{}]".format(", ".join(platform.categories[Category(x)] for x in sub.categories)) - if platform.enable_tag: - res += " {}".format(", ".join(sub.tags)) - else: - res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）" - if res[-1] != "\n": - res += "\n" + res = await generate_sub_list_text( + add_cookie_target_matcher, state, user_info, is_index=True, is_show_cookie=True + ) res += "请输入要关联 cookie 的订阅的序号\n输入'取消'中止" await MessageFactory(await parse_text(res)).send() @@ -64,7 +45,9 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): associated_cookie_ids = {cookie.id for cookie in associated_cookies} cookies = [cookie for cookie in cookies if cookie.id not in associated_cookie_ids] if not cookies: - await add_cookie_target_matcher.finish("当前平台暂无 Cookie，请使用“添加cookie”命令添加") + await add_cookie_target_matcher.finish( + "当前平台暂无可关联的 Cookie，请使用“添加cookie”命令添加或检查已关联的 Cookie" + ) state["cookies"] = cookies state["_prompt"] = "请选择一个 Cookie，已关联的 Cookie 不会显示\n" + "\n".join( [f"{idx}. {await get_cookie_friendly_name(cookie)}" for idx, cookie in enumerate(cookies, 1)] diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index d069496..71a06e3 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -6,11 +6,14 @@ from nonebot.adapters import Event from nonebot.typing import T_State from nonebot.matcher import Matcher from nonebot.permission import SUPERUSER -from nonebot_plugin_saa import extract_target from nonebot.params import Depends, EventToMe, EventPlainText +from nonebot_plugin_saa import PlatformTarget, extract_target +from ..config import config +from ..types import Category from ..platform import platform_manager from ..plugin_config import plugin_config +from ..apis import get_cookie_friendly_name def _configurable_to_me(to_me: bool = EventToMe()): @@ -60,3 +63,39 @@ def admin_permission(): permission = permission | GROUP_ADMIN | GROUP_OWNER return permission + + +async def generate_sub_list_text( + matcher: type[Matcher], state: T_State, user_info: PlatformTarget, is_index=False, is_show_cookie=False +): + sub_list = await config.list_subscribe(user_info) + if not sub_list: + await matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅") + res = "订阅的帐号为：\n" + state["sub_table"] = {} + for index, sub in enumerate(sub_list, 1): + state["sub_table"][index] = { + "platform_name": sub.target.platform_name, + "target": sub.target.target, + } + res += f"{index} " if is_index else "" + res += f"{sub.target.platform_name} {sub.target.target_name} {sub.target.target}\n" + if platform := platform_manager.get(sub.target.platform_name): + if platform.categories: + res += " [{}]".format(", ".join(platform.categories[Category(x)] for x in sub.categories)) + "\n" + if platform.enable_tag: + if sub.tags: + res += " {}".format(", ".join(sub.tags)) + "\n" + if is_show_cookie: + target_cookies = await config.get_cookie( + user=user_info, target=sub.target.target, platform_name=sub.target.platform_name + ) + if target_cookies: + res += " 关联的 Cookie：\n" + for cookie in target_cookies: + res += f" \t{await get_cookie_friendly_name(cookie)}\n" + + else: + res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）" + + return res From c828fd94e4d45b97a4942f1435efc6bb5aab2dd1 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sat, 31 Aug 2024 23:04:49 +0800 Subject: [PATCH 06/51] =?UTF-8?q?:recycle:=20cookie=20=E7=BB=84=E4=BB=B6?= =?UTF-8?q?=E4=B8=8D=E5=86=8D=E4=B8=8E=20user=20=E5=85=B3=E8=81=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../migrations/bdec8925b540_add_cookie.py | 63 ------------------- 1 file changed, 63 deletions(-) delete mode 100644 nonebot_bison/config/migrations/bdec8925b540_add_cookie.py diff --git a/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py b/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py deleted file mode 100644 index ac650f3..0000000 --- a/nonebot_bison/config/migrations/bdec8925b540_add_cookie.py +++ /dev/null @@ -1,63 +0,0 @@ -"""empty message - -Revision ID: bdec8925b540 -Revises: f9baef347cc8 -Create Date: 2024-08-22 20:53:08.850051 - -""" - -import sqlalchemy as sa -from alembic import op -from sqlalchemy import Text -from sqlalchemy.dialects import postgresql - -# revision identifiers, used by Alembic. -revision = "bdec8925b540" -down_revision = "f9baef347cc8" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.create_table( - "nonebot_bison_cookie", - sa.Column("id", sa.Integer(), nullable=False), - sa.Column("user_id", sa.Integer(), nullable=False), - sa.Column("platform_name", sa.String(length=20), nullable=False), - sa.Column("content", sa.String(length=1024), nullable=False), - sa.Column("last_usage", sa.DateTime(), nullable=False), - sa.Column("status", sa.String(length=20), nullable=False), - sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False), - sa.ForeignKeyConstraint( - ["user_id"], ["nonebot_bison_user.id"], name=op.f("fk_nonebot_bison_cookie_user_id_nonebot_bison_user") - ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")), - ) - op.create_table( - "nonebot_bison_cookietarget", - sa.Column("id", sa.Integer(), nullable=False), - sa.Column("target_id", sa.Integer(), nullable=False), - sa.Column("cookie_id", sa.Integer(), nullable=False), - sa.ForeignKeyConstraint( - ["cookie_id"], - ["nonebot_bison_cookie.id"], - name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"), - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["target_id"], - ["nonebot_bison_target.id"], - name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"), - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")), - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.drop_table("nonebot_bison_cookietarget") - op.drop_table("nonebot_bison_cookie") - # ### end Alembic commands ### From 1cd778c2e0506208b5520a755917c65f3692ea13 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sat, 31 Aug 2024 23:07:43 +0800 Subject: [PATCH 07/51] =?UTF-8?q?:recycle:=20cookie=20=E7=BB=84=E4=BB=B6?= =?UTF-8?q?=E4=B8=8D=E5=86=8D=E4=B8=8E=20user=20=E5=85=B3=E8=81=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 28 ++------- nonebot_bison/config/db_model.py | 3 - .../migrations/590dc2911ea7_add_cookie.py | 59 +++++++++++++++++++ nonebot_bison/sub_manager/add_cookie.py | 9 +-- .../sub_manager/add_cookie_target.py | 21 +++---- nonebot_bison/sub_manager/utils.py | 15 ++++- 6 files changed, 85 insertions(+), 50 deletions(-) create mode 100644 nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 97cc0db..909d6f9 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,14 +259,9 @@ class DBConfig: ) return res - async def get_cookie( - self, user: PlatformTarget = None, platform_name: str = None, target: T_Target = None - ) -> list[Cookie]: + async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]: async with create_session() as sess: - query = select(Cookie).distinct().join(User) - if user: - user_id = await sess.scalar(select(User.id).where(User.user_target == model_dump(user))) - query = query.where(Cookie.user_id == user_id) + query = select(Cookie).distinct() if platform_name: query = query.where(Cookie.platform_name == platform_name) query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) @@ -277,24 +272,9 @@ class DBConfig: res = [cookie for cookie in res if cookie.id in ids] return res - async def get_cookie_by_user_and_platform(self, user: PlatformTarget, platform_name: str) -> list[Cookie]: + async def add_cookie(self, platform_name: str, content: str) -> int: async with create_session() as sess: - res = await sess.scalar( - select(User) - .where(User.user_target == model_dump(user)) - .join(Cookie) - .where(Cookie.platform_name == platform_name) - .outerjoin(CookieTarget) - .options(selectinload(User.cookies)) - ) - if not res: - return [] - return res.cookies - - async def add_cookie(self, user: PlatformTarget, platform_name: str, content: str) -> int: - async with create_session() as sess: - user_obj = await sess.scalar(select(User).where(User.user_target == model_dump(user))) - cookie = Cookie(user=user_obj, platform_name=platform_name, content=content) + cookie = Cookie(platform_name=platform_name, content=content) sess.add(cookie) await sess.commit() await sess.refresh(cookie) diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 603ae21..2c74721 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -20,7 +20,6 @@ class User(Model): user_target: Mapped[dict] = mapped_column(JSON().with_variant(JSONB, "postgresql")) subscribes: Mapped[list["Subscribe"]] = relationship(back_populates="user") - cookies: Mapped[list["Cookie"]] = relationship(back_populates="user") @property def saa_target(self) -> PlatformTarget: @@ -73,14 +72,12 @@ class Subscribe(Model): class Cookie(Model): id: Mapped[int] = mapped_column(primary_key=True) - user_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_user.id")) 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)) status: Mapped[str] = mapped_column(String(20), default="") tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={}) - user: Mapped[User] = relationship(back_populates="cookies") targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie") diff --git a/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py b/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py new file mode 100644 index 0000000..c8fbd1d --- /dev/null +++ b/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py @@ -0,0 +1,59 @@ +"""empty message + +Revision ID: 590dc2911ea7 +Revises: f9baef347cc8 +Create Date: 2024-08-31 23:06:02.123932 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import Text +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "590dc2911ea7" +down_revision = "f9baef347cc8" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "nonebot_bison_cookie", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("platform_name", sa.String(length=20), nullable=False), + sa.Column("content", sa.String(length=1024), nullable=False), + sa.Column("last_usage", sa.DateTime(), nullable=False), + sa.Column("status", sa.String(length=20), nullable=False), + sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")), + ) + op.create_table( + "nonebot_bison_cookietarget", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("target_id", sa.Integer(), nullable=False), + sa.Column("cookie_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["cookie_id"], + ["nonebot_bison_cookie.id"], + name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["target_id"], + ["nonebot_bison_target.id"], + name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("nonebot_bison_cookietarget") + op.drop_table("nonebot_bison_cookie") + # ### end Alembic commands ### diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index f6ba356..8702a3c 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -1,21 +1,18 @@ from nonebot.typing import T_State from nonebot.matcher import Matcher from nonebot.params import Arg, ArgPlainText -from nonebot_plugin_saa import PlatformTarget from nonebot.adapters import Message, MessageTemplate from ..types import Target from ..config import config from ..platform import platform_manager from ..apis import check_sub_target_cookie -from .utils import common_platform, ensure_user_info, gen_handle_cancel +from .utils import common_platform, gen_handle_cancel def do_add_cookie(add_cookie: type[Matcher]): handle_cancel = gen_handle_cancel(add_cookie, "已中止添加cookie") - add_cookie.handle()(ensure_user_info(add_cookie)) - @add_cookie.handle() async def init_promote(state: T_State): state["_prompt"] = ( @@ -56,8 +53,8 @@ def do_add_cookie(add_cookie: type[Matcher]): state["name"] = await check_sub_target_cookie(state["platform"], Target(""), cookie_text) @add_cookie.handle() - async def add_cookie_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): - await config.add_cookie(user, state["platform"], state["cookie"]) + async def add_cookie_process(state: T_State): + await config.add_cookie(state["platform"], state["cookie"]) await add_cookie.finish( f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" ) diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index 7244736..5a1a4d0 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -1,25 +1,21 @@ from nonebot.typing import T_State from nonebot.matcher import Matcher -from nonebot.params import Arg, ArgPlainText +from nonebot.params import ArgPlainText +from nonebot_plugin_saa import MessageFactory from nonebot.internal.adapter import MessageTemplate -from nonebot_plugin_saa import MessageFactory, PlatformTarget from ..config import config from ..utils import parse_text from ..apis import get_cookie_friendly_name -from .utils import ensure_user_info, gen_handle_cancel, generate_sub_list_text +from .utils import gen_handle_cancel, generate_sub_list_text def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): handle_cancel = gen_handle_cancel(add_cookie_target_matcher, "已中止关联 cookie") - add_cookie_target_matcher.handle()(ensure_user_info(add_cookie_target_matcher)) - @add_cookie_target_matcher.handle() - async def init_promote(state: T_State, user_info: PlatformTarget = Arg("target_user_info")): - res = await generate_sub_list_text( - add_cookie_target_matcher, state, user_info, is_index=True, is_show_cookie=True - ) + async def init_promote(state: T_State): + res = await generate_sub_list_text(add_cookie_target_matcher, state, is_index=True, is_show_cookie=True) res += "请输入要关联 cookie 的订阅的序号\n输入'取消'中止" await MessageFactory(await parse_text(res)).send() @@ -34,11 +30,8 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def init_promote_cookie(state: T_State): - cookies = await config.get_cookie( - user=state["target_user_info"], platform_name=state["target"]["platform_name"] - ) + cookies = await config.get_cookie(platform_name=state["target"]["platform_name"]) associated_cookies = await config.get_cookie( - user=state["target_user_info"], target=state["target"]["target"], platform_name=state["target"]["platform_name"], ) @@ -62,7 +55,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): await add_cookie_target_matcher.reject("序号错误") @add_cookie_target_matcher.handle() - async def add_cookie_target_process(state: T_State, user: PlatformTarget = Arg("target_user_info")): + async def add_cookie_target_process(state: T_State): await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id) await add_cookie_target_matcher.finish( f"已关联 Cookie: {await get_cookie_friendly_name(state['cookie'])} " diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index 71a06e3..4b7dd6f 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -1,5 +1,7 @@ import contextlib from typing import Annotated +from itertools import groupby +from operator import attrgetter from nonebot.rule import Rule from nonebot.adapters import Event @@ -66,9 +68,16 @@ def admin_permission(): async def generate_sub_list_text( - matcher: type[Matcher], state: T_State, user_info: PlatformTarget, is_index=False, is_show_cookie=False + matcher: type[Matcher], state: T_State, user_info: PlatformTarget = None, is_index=False, is_show_cookie=False ): - sub_list = await config.list_subscribe(user_info) + if user_info: + sub_list = await config.list_subscribe(user_info) + else: + sub_list = await config.list_subs_with_all_info() + sub_list = [ + next(group) + for key, group in groupby(sorted(sub_list, key=attrgetter("target_id")), key=attrgetter("target_id")) + ] if not sub_list: await matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅") res = "订阅的帐号为：\n" @@ -88,7 +97,7 @@ async def generate_sub_list_text( res += " {}".format(", ".join(sub.tags)) + "\n" if is_show_cookie: target_cookies = await config.get_cookie( - user=user_info, target=sub.target.target, platform_name=sub.target.platform_name + target=sub.target.target, platform_name=sub.target.platform_name ) if target_cookies: res += " 关联的 Cookie：\n" From 7901b845ea2141fa4c8efa74baeb976413d9cd9a Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 2 Sep 2024 23:13:29 +0800 Subject: [PATCH 08/51] =?UTF-8?q?:sparkles:=20=E5=88=9D=E6=AD=A5=E5=AE=9E?= =?UTF-8?q?=E7=8E=B0=E6=90=BA=E5=B8=A6cookie=E8=AF=B7=E6=B1=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/platform/rss.py | 4 +++- nonebot_bison/utils/site.py | 42 +++++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/nonebot_bison/platform/rss.py b/nonebot_bison/platform/rss.py index 895b1e2..a94f3db 100644 --- a/nonebot_bison/platform/rss.py +++ b/nonebot_bison/platform/rss.py @@ -9,6 +9,7 @@ from bs4 import BeautifulSoup as bs from ..post import Post from .platform import NewMessage from ..types import Target, RawPost +from ..utils.site import create_cookie_client_manager from ..utils import Site, text_fletten, text_similarity @@ -16,6 +17,7 @@ class RssSite(Site): name = "rss" schedule_type = "interval" schedule_setting = {"seconds": 30} + client_mgr = create_cookie_client_manager("rss") class RssPost(Post): @@ -63,7 +65,7 @@ class Rss(NewMessage): return post.id async def get_sub_list(self, target: Target) -> list[RawPost]: - client = await self.ctx.get_client() + client = await self.ctx.get_client(target) res = await client.get(target, timeout=10.0) feed = feedparser.parse(res) entries = feed.entries diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 9a71d14..ca11f07 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -1,9 +1,12 @@ +import json from typing import Literal from abc import ABC, abstractmethod +import httpx from httpx import AsyncClient from ..types import Target +from ..config import config from .http import http_client @@ -35,6 +38,45 @@ class DefaultClientManager(ClientManager): pass +class CookieClientManager(ClientManager): + _platform_name: str + + async def _choose_cookie(self, target: Target) -> dict[str, str]: + if not target: + return {} + cookies = await config.get_cookie_by_target(target, self._platform_name) + if not cookies: + return {} + cookie = sorted(cookies, key=lambda x: x.last_usage, reverse=True)[0] + return json.loads(cookie.content) + + async def get_client(self, target: Target | None) -> AsyncClient: + client = http_client() + cookie = await self._choose_cookie(target) + cookies = httpx.Cookies() + if cookie: + cookies.update(cookie) + client.cookies = cookies + return client + + async def get_client_for_static(self) -> AsyncClient: + pass + + async def get_query_name_client(self) -> AsyncClient: + pass + + async def refresh_client(self): + pass + + +def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: + return type( + "CookieClientManager", + (CookieClientManager,), + {"_platform_name": platform_name}, + ) + + class Site: schedule_type: Literal["date", "interval", "cron"] schedule_setting: dict From 4ce6b85f7958ff5164b641cd76756eae076fa45d Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Tue, 3 Sep 2024 09:39:27 +0800 Subject: [PATCH 09/51] =?UTF-8?q?:sparkles:=20weibo=20=E5=B8=A6=20cookie?= =?UTF-8?q?=20mvp?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/platform/weibo.py | 8 ++++++-- nonebot_bison/utils/site.py | 4 ++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/nonebot_bison/platform/weibo.py b/nonebot_bison/platform/weibo.py index 7dfca5b..cadcce6 100644 --- a/nonebot_bison/platform/weibo.py +++ b/nonebot_bison/platform/weibo.py @@ -13,6 +13,7 @@ from bs4 import BeautifulSoup as bs from ..post import Post from .platform import NewMessage from ..utils import Site, http_client +from ..utils.site import create_cookie_client_manager from ..types import Tag, Target, RawPost, ApiError, Category _HEADER = { @@ -39,6 +40,7 @@ class WeiboSite(Site): name = "weibo.com" schedule_type = "interval" schedule_setting = {"seconds": 3} + client_mgr = create_cookie_client_manager("weibo") class Weibo(NewMessage): @@ -78,9 +80,11 @@ class Weibo(NewMessage): raise cls.ParseTargetException(prompt="正确格式:\n1. 用户数字UID\n2. https://weibo.com/u/xxxx") async def get_sub_list(self, target: Target) -> list[RawPost]: - client = await self.ctx.get_client() + client = await self.ctx.get_client(target) + header = {"Referer": f"https://m.weibo.cn/u/{target}", "MWeibo-Pwa": "1", "X-Requested-With": "XMLHttpRequest"} + # 获取 cookie 见 https://docs.rsshub.app/zh/deploy/config#%E5%BE%AE%E5%8D%9A params = {"containerid": "107603" + target} - res = await client.get("https://m.weibo.cn/api/container/getIndex?", params=params, timeout=4.0) + res = await client.get("https://m.weibo.cn/api/container/getIndex?", headers=header, params=params, timeout=4.0) res_data = json.loads(res.text) if not res_data["ok"] and res_data["msg"] != "这里还没有内容": raise ApiError(res.request.url) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index ca11f07..4fc4e51 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -60,10 +60,10 @@ class CookieClientManager(ClientManager): return client async def get_client_for_static(self) -> AsyncClient: - pass + return http_client() async def get_query_name_client(self) -> AsyncClient: - pass + return http_client() async def refresh_client(self): pass From 055ed6e02a6f422954d72f6cfedf13b6084ceab1 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Tue, 3 Sep 2024 10:52:20 +0800 Subject: [PATCH 10/51] =?UTF-8?q?:sparkles:=20=E4=BD=BF=E7=94=A8=E9=97=AD?= =?UTF-8?q?=E5=8C=85=E5=AE=9E=E7=8E=B0client=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/context.py | 3 ++- nonebot_bison/utils/site.py | 31 ++++++++++++++++++++++++++----- 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/nonebot_bison/utils/context.py b/nonebot_bison/utils/context.py index 7f5f168..5359969 100644 --- a/nonebot_bison/utils/context.py +++ b/nonebot_bison/utils/context.py @@ -22,8 +22,9 @@ class ProcessContext: async def _log_to_ctx(r: Response): self._log_response(r) + existing_hooks = client.event_hooks["response"] hooks = { - "response": [_log_to_ctx], + "response": [*existing_hooks, _log_to_ctx], } client.event_hooks = hooks diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 4fc4e51..e6d3db9 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -4,10 +4,12 @@ from abc import ABC, abstractmethod import httpx from httpx import AsyncClient +from nonebot.log import logger from ..types import Target from ..config import config from .http import http_client +from ..config.db_model import Cookie class ClientManager(ABC): @@ -41,22 +43,41 @@ class DefaultClientManager(ClientManager): class CookieClientManager(ClientManager): _platform_name: str - async def _choose_cookie(self, target: Target) -> dict[str, str]: + def _generate_hook(self, cookie: Cookie): + async def _response_hook(resp: httpx.Response): + if not cookie: + logger.debug(f"未携带bison cookie: {resp.request.url}") + if resp.status_code == 200: + logger.debug(f"请求成功: {cookie.id} {resp.request.url}") + + else: + logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") + + return _response_hook + + async def _choose_cookie(self, target: Target) -> Cookie: if not target: - return {} + return Cookie(content="{}") cookies = await config.get_cookie_by_target(target, self._platform_name) if not cookies: - return {} + return Cookie(content="{}") cookie = sorted(cookies, key=lambda x: x.last_usage, reverse=True)[0] - return json.loads(cookie.content) + return cookie async def get_client(self, target: Target | None) -> AsyncClient: client = http_client() cookie = await self._choose_cookie(target) + if cookie.content != "{}": + logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") + else: + logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie") + cookies = httpx.Cookies() if cookie: - cookies.update(cookie) + cookies.update(json.loads(cookie.content)) client.cookies = cookies + client._bison_cookie = cookie + client.event_hooks = {"response": [self._generate_hook(cookie)]} return client async def get_client_for_static(self) -> AsyncClient: From 498e7d60d411a23308649bbf4de093ee2fa454b3 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Tue, 3 Sep 2024 11:06:25 +0800 Subject: [PATCH 11/51] =?UTF-8?q?:sparkles:=20=E6=A0=B9=E6=8D=AE=E6=9D=A1?= =?UTF-8?q?=E4=BB=B6=E9=80=89=E6=8B=A9Cookie=20=E5=92=8C=20=E7=8A=B6?= =?UTF-8?q?=E6=80=81=E5=9B=9E=E5=86=99=E6=95=B0=E6=8D=AE=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/site.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index e6d3db9..8550b26 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -1,6 +1,7 @@ import json from typing import Literal from abc import ABC, abstractmethod +from datetime import datetime, timedelta import httpx from httpx import AsyncClient @@ -42,29 +43,41 @@ class DefaultClientManager(ClientManager): class CookieClientManager(ClientManager): _platform_name: str + _cookie_cd: int = 10 def _generate_hook(self, cookie: Cookie): + """hook函数生成器，用于回写请求状态到数据库""" + async def _response_hook(resp: httpx.Response): - if not cookie: + 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" else: logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") + cookie.status = "failed" + cookie.last_usage = datetime.now() + await config.update_cookie(cookie) return _response_hook 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() + ) if not cookies: return Cookie(content="{}") - cookie = sorted(cookies, key=lambda x: x.last_usage, reverse=True)[0] + cookie = max(cookies, key=lambda x: x.last_usage) return cookie async def get_client(self, target: Target | None) -> AsyncClient: + """获取 client，根据 target 选择 cookie""" client = http_client() cookie = await self._choose_cookie(target) if cookie.content != "{}": @@ -91,6 +104,7 @@ class CookieClientManager(ClientManager): def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: + """创建一个预定义为平台特化的 CookieClientManger""" return type( "CookieClientManager", (CookieClientManager,), From 01435eeded42ca86eb2c47f96def5be398fa19d2 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Wed, 4 Sep 2024 01:19:53 +0800 Subject: [PATCH 12/51] =?UTF-8?q?:sparkles:=20=E6=8F=90=E5=87=BAassemble?= =?UTF-8?q?=5Fclient=E6=96=B9=E6=B3=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/site.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 8550b26..913b815 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -85,6 +85,9 @@ class CookieClientManager(ClientManager): else: logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie") + return await self.assemble_client(client, cookie) + + async def assemble_client(self, client, cookie): cookies = httpx.Cookies() if cookie: cookies.update(json.loads(cookie.content)) @@ -104,7 +107,7 @@ class CookieClientManager(ClientManager): def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: - """创建一个预定义为平台特化的 CookieClientManger""" + """创建一个平台特化的 CookieClientManger""" return type( "CookieClientManager", (CookieClientManager,), From 3bd0867f0eb4b2cad452838c87fddb0242c900e7 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 5 Sep 2024 16:07:59 +0800 Subject: [PATCH 13/51] =?UTF-8?q?:sparkles:=20=E6=95=B0=E6=8D=AE=E5=BA=93C?= =?UTF-8?q?ookie=E8=A1=A8=E6=B7=BB=E5=8A=A0cd=E5=B1=9E=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_model.py | 1 + .../migrations/590dc2911ea7_add_cookie.py | 59 ------------------- 2 files changed, 1 insertion(+), 59 deletions(-) delete mode 100644 nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 2c74721..d7c84f9 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -76,6 +76,7 @@ class Cookie(Model): content: Mapped[str] = mapped_column(String(1024)) last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1)) status: Mapped[str] = mapped_column(String(20), default="") + cd: Mapped[int] = mapped_column(default=0) 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/config/migrations/590dc2911ea7_add_cookie.py b/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py deleted file mode 100644 index c8fbd1d..0000000 --- a/nonebot_bison/config/migrations/590dc2911ea7_add_cookie.py +++ /dev/null @@ -1,59 +0,0 @@ -"""empty message - -Revision ID: 590dc2911ea7 -Revises: f9baef347cc8 -Create Date: 2024-08-31 23:06:02.123932 - -""" - -import sqlalchemy as sa -from alembic import op -from sqlalchemy import Text -from sqlalchemy.dialects import postgresql - -# revision identifiers, used by Alembic. -revision = "590dc2911ea7" -down_revision = "f9baef347cc8" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.create_table( - "nonebot_bison_cookie", - sa.Column("id", sa.Integer(), nullable=False), - sa.Column("platform_name", sa.String(length=20), nullable=False), - sa.Column("content", sa.String(length=1024), nullable=False), - sa.Column("last_usage", sa.DateTime(), nullable=False), - sa.Column("status", sa.String(length=20), nullable=False), - sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False), - sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")), - ) - op.create_table( - "nonebot_bison_cookietarget", - sa.Column("id", sa.Integer(), nullable=False), - sa.Column("target_id", sa.Integer(), nullable=False), - sa.Column("cookie_id", sa.Integer(), nullable=False), - sa.ForeignKeyConstraint( - ["cookie_id"], - ["nonebot_bison_cookie.id"], - name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"), - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["target_id"], - ["nonebot_bison_target.id"], - name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"), - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")), - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - # ### commands auto generated by Alembic - please adjust! ### - op.drop_table("nonebot_bison_cookietarget") - op.drop_table("nonebot_bison_cookie") - # ### end Alembic commands ### From 370fc250f03de5d7375eca338c2c99ba3973cfcb Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 5 Sep 2024 19:32:47 +0800 Subject: [PATCH 14/51] =?UTF-8?q?:sparkles:=20=E6=95=B0=E6=8D=AE=E5=BA=93C?= =?UTF-8?q?ookie=E8=A1=A8=E6=B7=BB=E5=8A=A0is=5Funiversal=E5=B1=9E?= =?UTF-8?q?=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_model.py | 1 + .../migrations/fc2b8481bdde_add_cookie.py | 61 +++++++++++++++++++ nonebot_bison/utils/site.py | 9 ++- 3 files changed, 69 insertions(+), 2 deletions(-) create mode 100644 nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index d7c84f9..4362ee4 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -77,6 +77,7 @@ class Cookie(Model): last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1)) status: Mapped[str] = mapped_column(String(20), default="") cd: Mapped[int] = mapped_column(default=0) + 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/config/migrations/fc2b8481bdde_add_cookie.py b/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py new file mode 100644 index 0000000..b296ecf --- /dev/null +++ b/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py @@ -0,0 +1,61 @@ +"""empty message + +Revision ID: fc2b8481bdde +Revises: f9baef347cc8 +Create Date: 2024-09-05 19:31:59.366754 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import Text +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "fc2b8481bdde" +down_revision = "f9baef347cc8" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "nonebot_bison_cookie", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("platform_name", sa.String(length=20), nullable=False), + sa.Column("content", sa.String(length=1024), nullable=False), + sa.Column("last_usage", sa.DateTime(), nullable=False), + sa.Column("status", sa.String(length=20), nullable=False), + sa.Column("cd", sa.Integer(), nullable=False), + sa.Column("is_universal", sa.Boolean(), nullable=False), + sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")), + ) + op.create_table( + "nonebot_bison_cookietarget", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("target_id", sa.Integer(), nullable=False), + sa.Column("cookie_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["cookie_id"], + ["nonebot_bison_cookie.id"], + name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["target_id"], + ["nonebot_bison_target.id"], + name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("nonebot_bison_cookietarget") + op.drop_table("nonebot_bison_cookie") + # ### end Alembic commands ### diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 913b815..ac3ec94 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -71,11 +71,16 @@ class CookieClientManager(ClientManager): 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) 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() @@ -85,9 +90,9 @@ class CookieClientManager(ClientManager): else: logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie") - return await self.assemble_client(client, cookie) + return await self._assemble_client(client, cookie) - async def assemble_client(self, client, cookie): + async def _assemble_client(self, client, cookie): cookies = httpx.Cookies() if cookie: cookies.update(json.loads(cookie.content)) From cf3966e69b8f960c3f172a695e52a81e7ddc749a Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 00:14:35 +0800 Subject: [PATCH 15/51] =?UTF-8?q?:sparkles:=20=E6=B7=BB=E5=8A=A0cookie?= =?UTF-8?q?=E5=8F=AA=E6=98=BE=E7=A4=BA=E6=94=AF=E6=8C=81=E7=9A=84Platform?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/sub_manager/add_cookie.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 8702a3c..2b0211b 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -5,6 +5,7 @@ from nonebot.adapters import Message, MessageTemplate from ..types import Target from ..config import config +from ..config.db_model import Cookie from ..platform import platform_manager from ..apis import check_sub_target_cookie from .utils import common_platform, gen_handle_cancel @@ -18,7 +19,11 @@ def do_add_cookie(add_cookie: type[Matcher]): state["_prompt"] = ( "请输入想要添加 Cookie 的平台，目前支持，请输入冒号左边的名称：\n" + "".join( - [f"{platform_name}: {platform_manager[platform_name].name}\n" for platform_name in common_platform] + [ + f"{platform_name}: {platform_manager[platform_name].name}\n" + for platform_name in common_platform + if hasattr(platform_manager[platform_name].site.client_mgr, "_cookie_client_manger_") + ] ) + "要查看全部平台请输入：“全部”\n中止添加cookie过程请输入：“取消”" ) @@ -27,7 +32,11 @@ def do_add_cookie(add_cookie: type[Matcher]): async def parse_platform(state: T_State, platform: str = ArgPlainText()) -> None: if platform == "全部": message = "全部平台\n" + "\n".join( - [f"{platform_name}: {platform.name}" for platform_name, platform in platform_manager.items()] + [ + f"{platform_name}: {platform.name}" + for platform_name, platform in platform_manager.items() + if hasattr(platform_manager[platform_name].site.client_mgr, "_cookie_client_manger_") + ] ) await add_cookie.reject(message) elif platform == "取消": @@ -54,7 +63,9 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.handle() async def add_cookie_process(state: T_State): - await config.add_cookie(state["platform"], state["cookie"]) + cookie = Cookie(platform_name=state["platform"], content=state["cookie"]) + cookie = platform_manager[state["platform"]].site.init_cookie(cookie) + await config.add_cookie(cookie) await add_cookie.finish( f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" ) From 0ce2893911c20aa435b64bc77359977ecef13363 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 00:25:33 +0800 Subject: [PATCH 16/51] =?UTF-8?q?:sparkles:=20=E5=8C=BF=E5=90=8Dcookie?= =?UTF-8?q?=E5=92=8C=E7=94=A8=E6=88=B7cookie=E4=B8=80=E8=B5=B7=E8=B0=83?= =?UTF-8?q?=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) From 6537f01a34b07a07de6388c631148bea29e42a00 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 00:34:52 +0800 Subject: [PATCH 17/51] =?UTF-8?q?:sparkles:=20=E9=9B=86=E4=B8=AD=E5=88=A4?= =?UTF-8?q?=E6=96=AD=E6=98=AF=E5=90=A6=E4=B8=BACookieClientManager?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/scheduler/manager.py | 3 ++- nonebot_bison/sub_manager/add_cookie.py | 5 +++-- nonebot_bison/utils/site.py | 4 ++++ 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/nonebot_bison/scheduler/manager.py b/nonebot_bison/scheduler/manager.py index 06486eb..3c4a5a3 100644 --- a/nonebot_bison/scheduler/manager.py +++ b/nonebot_bison/scheduler/manager.py @@ -7,6 +7,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 scheduler_dict: dict[type[Site], Scheduler] = {} @@ -30,7 +31,7 @@ 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_"): + if is_cookie_client_manager(site.client_mgr): 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.") diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 2b0211b..41384da 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -8,6 +8,7 @@ from ..config import config from ..config.db_model import Cookie from ..platform import platform_manager from ..apis import check_sub_target_cookie +from ..utils.site import is_cookie_client_manager from .utils import common_platform, gen_handle_cancel @@ -22,7 +23,7 @@ def do_add_cookie(add_cookie: type[Matcher]): [ f"{platform_name}: {platform_manager[platform_name].name}\n" for platform_name in common_platform - if hasattr(platform_manager[platform_name].site.client_mgr, "_cookie_client_manger_") + if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) ] ) + "要查看全部平台请输入：“全部”\n中止添加cookie过程请输入：“取消”" @@ -35,7 +36,7 @@ def do_add_cookie(add_cookie: type[Matcher]): [ f"{platform_name}: {platform.name}" for platform_name, platform in platform_manager.items() - if hasattr(platform_manager[platform_name].site.client_mgr, "_cookie_client_manger_") + if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) ] ) await add_cookie.reject(message) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 509fee8..d7618d6 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -41,6 +41,10 @@ class DefaultClientManager(ClientManager): pass +def is_cookie_client_manager(manger: type[ClientManager]) -> bool: + return hasattr(manger, "_cookie_client_manger_") + + class CookieClientManager(ClientManager): _cookie_client_manger_ = True _platform_name: str From b61bde6e3fcd76018c1e0fbe4017503c7206e93a Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 00:39:57 +0800 Subject: [PATCH 18/51] =?UTF-8?q?:sparkles:=20=E5=85=B3=E8=81=94cookie?= =?UTF-8?q?=E6=97=B6=EF=BC=8C=E5=8F=AA=E6=98=BE=E7=A4=BA=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E7=9A=84=E8=AE=A2=E9=98=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/apis.py | 3 ++- nonebot_bison/sub_manager/add_cookie_target.py | 4 +++- nonebot_bison/sub_manager/utils.py | 14 +++++++++++++- 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/nonebot_bison/apis.py b/nonebot_bison/apis.py index c4f247a..fb44654 100644 --- a/nonebot_bison/apis.py +++ b/nonebot_bison/apis.py @@ -1,4 +1,5 @@ from .types import Target +from .utils import text_fletten from .config.db_model import Cookie from .scheduler import scheduler_dict from .platform import platform_manager @@ -20,4 +21,4 @@ async def check_sub_target_cookie(platform_name: str, target: Target, cookie: st async def get_cookie_friendly_name(cookie: Cookie): # TODO - return f"{cookie.platform_name} [{cookie.content[:10]}]" + return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index 5a1a4d0..be3fe89 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -15,7 +15,9 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def init_promote(state: T_State): - res = await generate_sub_list_text(add_cookie_target_matcher, state, is_index=True, is_show_cookie=True) + res = await generate_sub_list_text( + add_cookie_target_matcher, state, is_index=True, is_show_cookie=True, is_hide_no_cookie_platfrom=True + ) res += "请输入要关联 cookie 的订阅的序号\n输入'取消'中止" await MessageFactory(await parse_text(res)).send() diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index 4b7dd6f..8edf7f8 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -16,6 +16,7 @@ from ..types import Category from ..platform import platform_manager from ..plugin_config import plugin_config from ..apis import get_cookie_friendly_name +from ..utils.site import is_cookie_client_manager def _configurable_to_me(to_me: bool = EventToMe()): @@ -68,7 +69,12 @@ def admin_permission(): async def generate_sub_list_text( - matcher: type[Matcher], state: T_State, user_info: PlatformTarget = None, is_index=False, is_show_cookie=False + matcher: type[Matcher], + state: T_State, + user_info: PlatformTarget = None, + is_index=False, + is_show_cookie=False, + is_hide_no_cookie_platfrom=False, ): if user_info: sub_list = await config.list_subscribe(user_info) @@ -78,6 +84,12 @@ async def generate_sub_list_text( next(group) for key, group in groupby(sorted(sub_list, key=attrgetter("target_id")), key=attrgetter("target_id")) ] + if is_hide_no_cookie_platfrom: + sub_list = [ + sub + for sub in sub_list + if is_cookie_client_manager(platform_manager.get(sub.target.platform_name).site.client_mgr) + ] if not sub_list: await matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅") res = "订阅的帐号为：\n" From afd1bee7627c0f35a7804c727465e6b92c8c5de0 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 01:08:32 +0800 Subject: [PATCH 19/51] =?UTF-8?q?:sparkles:=20=E5=B0=86get=5Fcookie=5Ffrie?= =?UTF-8?q?ndly=5Fname=E5=92=8Cvalid=5Fcookie=E7=A7=BB=E5=8A=A8=E5=88=B0cc?= =?UTF-8?q?m=E5=86=85=E9=83=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/apis.py | 12 ----- nonebot_bison/sub_manager/add_cookie.py | 21 ++++----- .../sub_manager/add_cookie_target.py | 13 +++-- nonebot_bison/sub_manager/utils.py | 8 ++-- nonebot_bison/utils/site.py | 47 +++++++++++++------ 5 files changed, 56 insertions(+), 45 deletions(-) diff --git a/nonebot_bison/apis.py b/nonebot_bison/apis.py index fb44654..f13580b 100644 --- a/nonebot_bison/apis.py +++ b/nonebot_bison/apis.py @@ -1,6 +1,4 @@ from .types import Target -from .utils import text_fletten -from .config.db_model import Cookie from .scheduler import scheduler_dict from .platform import platform_manager @@ -12,13 +10,3 @@ async def check_sub_target(platform_name: str, target: Target): client = await scheduler.client_mgr.get_query_name_client() return await platform_manager[platform_name].get_target_name(client, target) - - -async def check_sub_target_cookie(platform_name: str, target: Target, cookie: str): - # TODO - return "check pass" - - -async def get_cookie_friendly_name(cookie: Cookie): - # TODO - return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 41384da..c10fc27 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -1,15 +1,15 @@ +from typing import cast + from nonebot.typing import T_State from nonebot.matcher import Matcher from nonebot.params import Arg, ArgPlainText from nonebot.adapters import Message, MessageTemplate -from ..types import Target from ..config import config from ..config.db_model import Cookie from ..platform import platform_manager -from ..apis import check_sub_target_cookie -from ..utils.site import is_cookie_client_manager from .utils import common_platform, gen_handle_cancel +from ..utils.site import CookieClientManager, is_cookie_client_manager def do_add_cookie(add_cookie: type[Matcher]): @@ -48,24 +48,21 @@ def do_add_cookie(add_cookie: type[Matcher]): await add_cookie.reject("平台输入错误") @add_cookie.handle() - async def prepare_get_id(matcher: Matcher, state: T_State): - cur_platform = platform_manager[state["platform"]] - if cur_platform.has_target: - state["_prompt"] = "请输入 Cookie" - else: - matcher.set_arg("cookie", None) # type: ignore - state["id"] = "default" + async def prepare_get_id(state: T_State): + state["_prompt"] = "请输入 Cookie" @add_cookie.got("cookie", MessageTemplate("{_prompt}"), [handle_cancel]) async def got_cookie(state: T_State, cookie: Message = Arg()): + client_mgr: CookieClientManager = platform_manager[state["platform"]].site.client_mgr cookie_text = cookie.extract_plain_text() state["cookie"] = cookie_text - state["name"] = await check_sub_target_cookie(state["platform"], Target(""), cookie_text) + state["name"] = await client_mgr.valid_cookie(cookie_text) @add_cookie.handle() async def add_cookie_process(state: T_State): cookie = Cookie(platform_name=state["platform"], content=state["cookie"]) - cookie = platform_manager[state["platform"]].site.init_cookie(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 add_cookie.finish( f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index be3fe89..1650428 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -1,3 +1,5 @@ +from typing import cast + from nonebot.typing import T_State from nonebot.matcher import Matcher from nonebot.params import ArgPlainText @@ -6,7 +8,8 @@ from nonebot.internal.adapter import MessageTemplate from ..config import config from ..utils import parse_text -from ..apis import get_cookie_friendly_name +from ..platform import platform_manager +from ..utils.site import CookieClientManager from .utils import gen_handle_cancel, generate_sub_list_text @@ -44,8 +47,10 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): "当前平台暂无可关联的 Cookie，请使用“添加cookie”命令添加或检查已关联的 Cookie" ) state["cookies"] = cookies + + client_mgr = cast(CookieClientManager, platform_manager[cookies[0].platform_name].site.client_mgr) state["_prompt"] = "请选择一个 Cookie，已关联的 Cookie 不会显示\n" + "\n".join( - [f"{idx}. {await get_cookie_friendly_name(cookie)}" for idx, cookie in enumerate(cookies, 1)] + [f"{idx}. {await client_mgr.get_cookie_friendly_name(cookie)}" for idx, cookie in enumerate(cookies, 1)] ) @add_cookie_target_matcher.got("cookie_idx", MessageTemplate("{_prompt}"), [handle_cancel]) @@ -59,7 +64,9 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def add_cookie_target_process(state: T_State): await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id) + cookie = state["cookie"] + client_mgr = cast(CookieClientManager, platform_manager[cookie.platform_name].site.client_mgr) await add_cookie_target_matcher.finish( - f"已关联 Cookie: {await get_cookie_friendly_name(state['cookie'])} " + f"已关联 Cookie: {await client_mgr.get_cookie_friendly_name(cookie)} " f"到订阅 {state['target']['platform_name']} {state['target']['target']}" ) diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index 8edf7f8..96b080c 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -1,7 +1,7 @@ import contextlib -from typing import Annotated from itertools import groupby from operator import attrgetter +from typing import Annotated, cast from nonebot.rule import Rule from nonebot.adapters import Event @@ -15,8 +15,7 @@ from ..config import config from ..types import Category from ..platform import platform_manager from ..plugin_config import plugin_config -from ..apis import get_cookie_friendly_name -from ..utils.site import is_cookie_client_manager +from ..utils.site import CookieClientManager, is_cookie_client_manager def _configurable_to_me(to_me: bool = EventToMe()): @@ -114,7 +113,8 @@ async def generate_sub_list_text( if target_cookies: res += " 关联的 Cookie：\n" for cookie in target_cookies: - res += f" \t{await get_cookie_friendly_name(cookie)}\n" + client_mgr = cast(CookieClientManager, platform_manager[cookie.platform_name].site.client_mgr) + res += f" \t{await client_mgr.get_cookie_friendly_name(cookie)}\n" else: res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）" diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index d7618d6..e241f49 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -50,6 +50,39 @@ class CookieClientManager(ClientManager): _platform_name: str _cookie_cd: int = 10 + @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时使用""" + cookie.cd = cls._cookie_cd + return cookie + + @classmethod + async def valid_cookie(cls, content: str) -> bool: + """验证cookie是否有效，添加cookie时考用，可根据平台的具体情况进行重写""" + try: + data = json.loads(content) + if not isinstance(data, dict): + raise ValueError + except Exception: + return False + return True + + @classmethod + async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: + """获取友好的cookie名字，用于展示""" + from . import text_fletten + + return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") + def _generate_hook(self, cookie: Cookie): """hook函数生成器，用于回写请求状态到数据库""" @@ -65,20 +98,6 @@ 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: From 418a941448ad133fff19e7c530281bdbbc8437be Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 10:08:08 +0800 Subject: [PATCH 20/51] =?UTF-8?q?:sparkles:=20=E6=B7=BB=E5=8A=A0=E4=B8=8D?= =?UTF-8?q?=E5=90=88=E6=B3=95cookie=E7=9A=84=E6=8F=90=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/sub_manager/add_cookie.py | 3 ++- nonebot_bison/utils/site.py | 14 ++++++++------ 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index c10fc27..9ae36e3 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -55,8 +55,9 @@ def do_add_cookie(add_cookie: type[Matcher]): async def got_cookie(state: T_State, cookie: Message = Arg()): client_mgr: CookieClientManager = platform_manager[state["platform"]].site.client_mgr cookie_text = cookie.extract_plain_text() + if not await client_mgr.valid_cookie(cookie_text): + await add_cookie.reject("无效的 Cookie，请检查后重新输入，详情见<待添加的文档>") state["cookie"] = cookie_text - state["name"] = await client_mgr.valid_cookie(cookie_text) @add_cookie.handle() async def add_cookie_process(state: T_State): diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index e241f49..4aeb8de 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -1,5 +1,6 @@ import json from typing import Literal +from json import JSONDecodeError from abc import ABC, abstractmethod from datetime import datetime, timedelta @@ -61,30 +62,31 @@ class CookieClientManager(ClientManager): @classmethod async def init_cookie(cls, cookie: Cookie) -> Cookie: - """初始化cookie，添加用户cookie时使用""" + """初始化 cookie，添加用户 cookie 时使用""" cookie.cd = cls._cookie_cd + cookie.last_usage = datetime.now() # 使得优先使用用户 cookie return cookie @classmethod async def valid_cookie(cls, content: str) -> bool: - """验证cookie是否有效，添加cookie时考用，可根据平台的具体情况进行重写""" + """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" try: data = json.loads(content) if not isinstance(data, dict): - raise ValueError - except Exception: + return False + except JSONDecodeError: return False return True @classmethod async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: - """获取友好的cookie名字，用于展示""" + """获取 cookie 的友好名字，用于展示""" from . import text_fletten return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") def _generate_hook(self, cookie: Cookie): - """hook函数生成器，用于回写请求状态到数据库""" + """hook 函数生成器，用于回写请求状态到数据库""" async def _response_hook(resp: httpx.Response): if resp.status_code == 200: From 06079b98f7f1471556ec77de680098db92a55633 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 10:16:24 +0800 Subject: [PATCH 21/51] =?UTF-8?q?:sparkles:=20=E5=85=B3=E8=81=94cookie?= =?UTF-8?q?=E6=98=AF=E4=B8=8D=E6=98=BE=E7=A4=BA=E5=8C=BF=E5=90=8Dcookie?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 8 ++------ nonebot_bison/utils/site.py | 3 ++- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index fc5ea53..63a18c7 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -260,8 +260,9 @@ class DBConfig: return res async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]: + """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" async with create_session() as sess: - query = select(Cookie).distinct() + query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 if platform_name: query = query.where(Cookie.platform_name == platform_name) query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) @@ -277,12 +278,7 @@ class DBConfig: 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: diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 4aeb8de..56f7132 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -55,9 +55,10 @@ class CookieClientManager(ClientManager): async def init_universal_cookie(cls): """移除已有的匿名cookie，添加一个新的匿名cookie""" universal_cookies = await config.get_unviersal_cookie(cls._platform_name) + universal_cookie = Cookie(platform_name=cls._platform_name, content="{}", is_universal=True) for cookie in universal_cookies: await config.delete_cookie(cookie.id) - universal_cookie = Cookie(platform_name=cls._platform_name, content="{}", is_universal=True) + universal_cookie.id = cookie.id # 保持原有的id await config.add_cookie(universal_cookie) @classmethod From eddd3e42a15b1ebd0c86bcdb3ef0de425c8e983c Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 6 Sep 2024 11:33:39 +0800 Subject: [PATCH 22/51] =?UTF-8?q?:sparkles:=20=E4=BF=AE=E6=94=B9=5Fchoose?= =?UTF-8?q?=5Fcookie=E7=9A=84=E9=80=BB=E8=BE=91=E4=BB=A5=E6=94=AF=E6=8C=81?= =?UTF-8?q?no=20target=E7=9A=84Platform?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/site.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 56f7132..752dc3c 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -86,7 +86,7 @@ class CookieClientManager(ClientManager): return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") - def _generate_hook(self, cookie: Cookie): + def _generate_hook(self, cookie: Cookie) -> callable: """hook 函数生成器，用于回写请求状态到数据库""" async def _response_hook(resp: httpx.Response): @@ -103,12 +103,9 @@ class CookieClientManager(ClientManager): 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)), - *(await config.get_universal_cookie(self._platform_name)), - ] + cookies = await config.get_universal_cookie(self._platform_name) + if target: + cookies += await config.get_cookie(self._platform_name, target) 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 @@ -128,7 +125,8 @@ class CookieClientManager(ClientManager): return await self._assemble_client(client, cookie) - async def _assemble_client(self, client, cookie): + async def _assemble_client(self, client, cookie) -> AsyncClient: + """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" cookies = httpx.Cookies() if cookie: cookies.update(json.loads(cookie.content)) From bbc54921939e753edbcebd5ab66b083764785a62 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 12:22:06 +0800 Subject: [PATCH 23/51] =?UTF-8?q?:sparkles:=20=E4=B8=BA=E5=8C=BF=E5=90=8Dc?= =?UTF-8?q?ookie=E8=AE=BE=E7=BD=AE=E6=A0=87=E7=AD=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/site.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 752dc3c..95dd4a2 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -55,8 +55,12 @@ class CookieClientManager(ClientManager): async def init_universal_cookie(cls): """移除已有的匿名cookie，添加一个新的匿名cookie""" universal_cookies = await config.get_unviersal_cookie(cls._platform_name) - universal_cookie = Cookie(platform_name=cls._platform_name, content="{}", is_universal=True) + universal_cookie = Cookie( + platform_name=cls._platform_name, content="{}", is_universal=True, tags={"temporary": True} + ) for cookie in universal_cookies: + if not cookie.tags.get("temporary"): + continue await config.delete_cookie(cookie.id) universal_cookie.id = cookie.id # 保持原有的id await config.add_cookie(universal_cookie) From a6227828e36f793beb2d0614ab25ab69ad33ce96 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 12:59:44 +0800 Subject: [PATCH 24/51] =?UTF-8?q?:recycle:=20=E6=95=B4=E7=90=86=E4=BB=A3?= =?UTF-8?q?=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/sub_manager/del_cookie.py | 0 nonebot_bison/sub_manager/del_cookie_target.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 nonebot_bison/sub_manager/del_cookie.py create mode 100644 nonebot_bison/sub_manager/del_cookie_target.py diff --git a/nonebot_bison/sub_manager/del_cookie.py b/nonebot_bison/sub_manager/del_cookie.py new file mode 100644 index 0000000..e69de29 diff --git a/nonebot_bison/sub_manager/del_cookie_target.py b/nonebot_bison/sub_manager/del_cookie_target.py new file mode 100644 index 0000000..e69de29 From 61dcf879ce3d20aac55b1748de9378e3c0ee8638 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 13:03:47 +0800 Subject: [PATCH 25/51] =?UTF-8?q?:recycle:=20=E6=95=B4=E7=90=86=E4=BB=A3?= =?UTF-8?q?=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 17 +++++++----- nonebot_bison/sub_manager/__init__.py | 27 ++++++++++++++++--- nonebot_bison/sub_manager/del_cookie.py | 9 +++++++ .../sub_manager/del_cookie_target.py | 9 +++++++ 4 files changed, 52 insertions(+), 10 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 63a18c7..0a47e9f 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,7 +259,7 @@ class DBConfig: ) return res - async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]: + async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" async with create_session() as sess: query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 @@ -273,7 +273,7 @@ class DBConfig: res = [cookie for cookie in res if cookie.id in ids] return res - async def get_unviersal_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]: + async def get_unviersal_cookie(self, platform_name: str = None) -> Sequence[Cookie]: async with create_session() as sess: query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 if platform_name: @@ -312,7 +312,7 @@ class DBConfig: await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) await sess.commit() - async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> list[Cookie]: + async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: async with create_session() as sess: query = ( select(Cookie) @@ -320,16 +320,16 @@ class DBConfig: .join(Target) .where(Target.platform_name == platform_name, Target.target == target) ) - return list((await sess.scalars(query)).all()) + return (await sess.scalars(query)).all() - async def get_universal_cookie(self, platform_name: str) -> list[Cookie]: + async def get_universal_cookie(self, platform_name: str) -> Sequence[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()) + return (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: @@ -358,5 +358,10 @@ class DBConfig: ) await sess.commit() + async def get_cookie_target(self) -> Sequence[CookieTarget]: + async with create_session() as sess: + query = select(CookieTarget).outerjoin(Target).options(selectinload(CookieTarget.target)) + return (await sess.scalars(query)).all() + config = DBConfig() diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index b16dcf7..943efd1 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -15,7 +15,9 @@ from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub from .add_cookie import do_add_cookie +from .del_cookie import do_del_cookie from .add_cookie_target import do_add_cookie_target +from .del_cookie_target import do_del_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info add_sub_matcher = on_command( @@ -45,23 +47,40 @@ do_del_sub(del_sub_matcher) add_cookie_matcher = on_command( "添加cookie", rule=configurable_to_me, - permission=admin_permission(), + permission=SUPERUSER, priority=5, block=True, ) -add_cookie_matcher.handle()(set_target_user_info) do_add_cookie(add_cookie_matcher) add_cookie_target_matcher = on_command( "关联cookie", rule=configurable_to_me, - permission=admin_permission(), + permission=SUPERUSER, priority=5, block=True, ) -add_cookie_target_matcher.handle()(set_target_user_info) do_add_cookie_target(add_cookie_target_matcher) +del_cookie_target_matcher = on_command( + "取消关联cookie", + rule=configurable_to_me, + permission=SUPERUSER, + priority=5, + block=True, +) +do_del_cookie_target(del_cookie_target_matcher) + +del_cookie_matcher = on_command( + "删除cookie", + rule=configurable_to_me, + permission=SUPERUSER, + priority=5, + block=True, +) +do_del_cookie(del_cookie_matcher) + + group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) group_handle_cancel = gen_handle_cancel(group_manage_matcher, "已取消") diff --git a/nonebot_bison/sub_manager/del_cookie.py b/nonebot_bison/sub_manager/del_cookie.py index e69de29..39d59b8 100644 --- a/nonebot_bison/sub_manager/del_cookie.py +++ b/nonebot_bison/sub_manager/del_cookie.py @@ -0,0 +1,9 @@ +from nonebot.matcher import Matcher + +from .utils import ensure_user_info, gen_handle_cancel + + +def do_del_cookie(del_cookie: type[Matcher]): + handle_cancel = gen_handle_cancel(del_cookie, "删除中止") + + del_cookie.handle()(ensure_user_info(del_cookie)) diff --git a/nonebot_bison/sub_manager/del_cookie_target.py b/nonebot_bison/sub_manager/del_cookie_target.py index e69de29..4187fe4 100644 --- a/nonebot_bison/sub_manager/del_cookie_target.py +++ b/nonebot_bison/sub_manager/del_cookie_target.py @@ -0,0 +1,9 @@ +from nonebot.matcher import Matcher + +from .utils import ensure_user_info, gen_handle_cancel + + +def do_del_cookie_target(del_cookie_target: type[Matcher]): + handle_cancel = gen_handle_cancel(del_cookie_target, "删除中止") + + del_cookie_target.handle()(ensure_user_info(del_cookie_target)) From 940301a6fc06df5d5d3a0d8fbf79f147b7df2a67 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 15:17:19 +0800 Subject: [PATCH 26/51] =?UTF-8?q?:sparkles:=20=E5=8F=96=E6=B6=88=E5=85=B3?= =?UTF-8?q?=E8=81=94cookie=20=E5=AF=B9=E8=AF=9D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 19 +++++-- .../sub_manager/del_cookie_target.py | 50 +++++++++++++++++-- 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 0a47e9f..c566187 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -358,10 +358,23 @@ class DBConfig: ) await sess.commit() - async def get_cookie_target(self) -> Sequence[CookieTarget]: + async def delete_cookie_target_by_id(self, cookie_target_id: int): async with create_session() as sess: - query = select(CookieTarget).outerjoin(Target).options(selectinload(CookieTarget.target)) - return (await sess.scalars(query)).all() + await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) + await sess.commit() + + async def get_cookie_target(self) -> list[CookieTarget]: + async with create_session() as sess: + query = ( + select(CookieTarget) + .outerjoin(Target) + .options(selectinload(CookieTarget.target)) + .outerjoin(Cookie) + .options(selectinload(CookieTarget.cookie)) + ) + res = list((await sess.scalars(query)).all()) + res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) + return res config = DBConfig() diff --git a/nonebot_bison/sub_manager/del_cookie_target.py b/nonebot_bison/sub_manager/del_cookie_target.py index 4187fe4..d161924 100644 --- a/nonebot_bison/sub_manager/del_cookie_target.py +++ b/nonebot_bison/sub_manager/del_cookie_target.py @@ -1,9 +1,51 @@ -from nonebot.matcher import Matcher +from typing import cast -from .utils import ensure_user_info, gen_handle_cancel +from nonebot.typing import T_State +from nonebot.matcher import Matcher +from nonebot.params import EventPlainText +from nonebot_plugin_saa import MessageFactory + +from ..config import config +from ..utils import parse_text +from .utils import gen_handle_cancel +from ..platform import platform_manager +from ..utils.site import CookieClientManager def do_del_cookie_target(del_cookie_target: type[Matcher]): - handle_cancel = gen_handle_cancel(del_cookie_target, "删除中止") + handle_cancel = gen_handle_cancel(del_cookie_target, "取消关联中止") - del_cookie_target.handle()(ensure_user_info(del_cookie_target)) + @del_cookie_target.handle() + async def send_list(state: T_State): + cookie_targets = await config.get_cookie_target() + if not cookie_targets: + await del_cookie_target.finish("暂无已关联 Cookie\n请使用“添加cookie”命令添加关联") + res = "已关联的 Cookie 为：\n" + state["cookie_target_table"] = {} + for index, cookie_target in enumerate(cookie_targets, 1): + client_mgr = cast(CookieClientManager, platform_manager[cookie_target.target.platform_name].site.client_mgr) + friendly_name = await client_mgr.get_cookie_friendly_name(cookie_target.cookie) + state["cookie_target_table"][index] = { + "platform_name": cookie_target.target.platform_name, + "target": cookie_target.target, + "friendly_name": friendly_name, + "cookie_target": cookie_target, + } + res += f"{index} {cookie_target.target.platform_name} {cookie_target.target.target_name} {friendly_name}\n" + if res[-1] != "\n": + res += "\n" + res += "请输入要删除的关联的序号\n输入'取消'中止" + await MessageFactory(await parse_text(res)).send() + + @del_cookie_target.receive(parameterless=[handle_cancel]) + async def do_del( + state: T_State, + index_str: str = EventPlainText(), + ): + try: + index = int(index_str) + await config.delete_cookie_target_by_id(state["cookie_target_table"][index]["cookie_target"].id) + except Exception: + await del_cookie_target.reject("删除错误") + else: + await del_cookie_target.finish("删除成功") From 7c9e191f4000290c112eded61786bdbee254aea1 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 15:56:44 +0800 Subject: [PATCH 27/51] =?UTF-8?q?:sparkles:=20=E5=88=A0=E9=99=A4cookie=20?= =?UTF-8?q?=E5=AF=B9=E8=AF=9D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 10 +++++- nonebot_bison/sub_manager/del_cookie.py | 41 +++++++++++++++++++++++-- nonebot_bison/utils/site.py | 2 +- tests/config/test_cookie.py | 2 +- 4 files changed, 50 insertions(+), 5 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index c566187..5cb0152 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -307,8 +307,16 @@ class DBConfig: cookie_in_db.tags = cookie.tags await sess.commit() - async def delete_cookie(self, cookie_id: int): + async def delete_cookie_by_id(self, cookie_id: int): async with create_session() as sess: + cookie = await sess.scalar( + select(Cookie) + .where(Cookie.id == cookie_id) + .outerjoin(CookieTarget) + .options(selectinload(Cookie.targets)) + ) + if len(cookie.targets) > 0: + raise Exception(f"cookie {cookie.id} in use") await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) await sess.commit() diff --git a/nonebot_bison/sub_manager/del_cookie.py b/nonebot_bison/sub_manager/del_cookie.py index 39d59b8..4054078 100644 --- a/nonebot_bison/sub_manager/del_cookie.py +++ b/nonebot_bison/sub_manager/del_cookie.py @@ -1,9 +1,46 @@ +from nonebot.typing import T_State from nonebot.matcher import Matcher +from nonebot.params import EventPlainText +from nonebot_plugin_saa import MessageFactory -from .utils import ensure_user_info, gen_handle_cancel +from ..config import config +from ..utils import parse_text +from .utils import gen_handle_cancel +from ..platform import platform_manager def do_del_cookie(del_cookie: type[Matcher]): handle_cancel = gen_handle_cancel(del_cookie, "删除中止") - del_cookie.handle()(ensure_user_info(del_cookie)) + @del_cookie.handle() + async def send_list(state: T_State): + cookies = await config.get_cookie() + if not cookies: + await del_cookie.finish("暂无已添加 Cookie\n请使用“添加cookie”命令添加") + res = "已添加的 Cookie 为：\n" + state["cookie_table"] = {} + for index, cookie in enumerate(cookies, 1): + state["cookie_table"][index] = cookie + client_mgr = platform_manager[cookie.platform_name].site.client_mgr + friendly_name = await client_mgr.get_cookie_friendly_name(cookie) + res += f"{index} {cookie.platform_name} {friendly_name} {len(cookie.targets)}个关联\n" + if res[-1] != "\n": + res += "\n" + res += "请输入要删除的 Cookie 的序号\n输入'取消'中止" + await MessageFactory(await parse_text(res)).send() + + @del_cookie.receive(parameterless=[handle_cancel]) + async def do_del( + state: T_State, + index_str: str = EventPlainText(), + ): + try: + index = int(index_str) + cookie = state["cookie_table"][index] + if cookie.targets: + await del_cookie.reject("只能删除未关联的 Cookie，请使用“取消关联cookie”命令取消关联") + await config.delete_cookie_by_id(cookie.id) + except Exception: + await del_cookie.reject("删除错误") + else: + await del_cookie.finish("删除成功") diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 95dd4a2..b53aab2 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -61,7 +61,7 @@ class CookieClientManager(ClientManager): for cookie in universal_cookies: if not cookie.tags.get("temporary"): continue - await config.delete_cookie(cookie.id) + await config.delete_cookie_by_id(cookie.id) universal_cookie.id = cookie.id # 保持原有的id await config.add_cookie(universal_cookie) diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py index 683b104..b13fa3f 100644 --- a/tests/config/test_cookie.py +++ b/tests/config/test_cookie.py @@ -73,7 +73,7 @@ async def test_cookie_by_user(app: App, init_scheduler): assert cookies[0].status == cookie.status assert cookies[0].tags == cookie.tags - await config.delete_cookie(cookies[0].id) + await config.delete_cookie_by_id(cookies[0].id) cookies = await config.get_cookie(TargetQQGroup(group_id=123)) assert len(cookies) == 0 From ce1f1bbedbc8c99a5e2d65f5041c74115765687a Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 18:12:00 +0800 Subject: [PATCH 28/51] =?UTF-8?q?:sparkles:=20=E5=8F=88=E6=9D=A5=E6=94=B9?= =?UTF-8?q?=E6=95=B0=E6=8D=AE=E5=BA=93=E4=BA=86(?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_model.py | 2 +- ...b8481bdde_add_cookie.py => be215495b122_add_cookie.py} | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) rename nonebot_bison/config/migrations/{fc2b8481bdde_add_cookie.py => be215495b122_add_cookie.py} (92%) diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index 590f978..fcd33bc 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -72,7 +72,7 @@ class Subscribe(Model): class Cookie(Model): id: Mapped[int] = mapped_column(primary_key=True) - platform_name: Mapped[str] = mapped_column(String(20)) + site_name: Mapped[str] = mapped_column(String(100)) content: Mapped[str] = mapped_column(String(1024)) # 最后使用的时刻 last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1)) diff --git a/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py b/nonebot_bison/config/migrations/be215495b122_add_cookie.py similarity index 92% rename from nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py rename to nonebot_bison/config/migrations/be215495b122_add_cookie.py index b296ecf..87e7bb1 100644 --- a/nonebot_bison/config/migrations/fc2b8481bdde_add_cookie.py +++ b/nonebot_bison/config/migrations/be215495b122_add_cookie.py @@ -1,8 +1,8 @@ """empty message -Revision ID: fc2b8481bdde +Revision ID: be215495b122 Revises: f9baef347cc8 -Create Date: 2024-09-05 19:31:59.366754 +Create Date: 2024-09-08 18:12:43.540818 """ @@ -12,7 +12,7 @@ from sqlalchemy import Text from sqlalchemy.dialects import postgresql # revision identifiers, used by Alembic. -revision = "fc2b8481bdde" +revision = "be215495b122" down_revision = "f9baef347cc8" branch_labels = None depends_on = None @@ -23,7 +23,7 @@ def upgrade() -> None: op.create_table( "nonebot_bison_cookie", sa.Column("id", sa.Integer(), nullable=False), - sa.Column("platform_name", sa.String(length=20), nullable=False), + sa.Column("site_name", sa.String(length=100), nullable=False), sa.Column("content", sa.String(length=1024), nullable=False), sa.Column("last_usage", sa.DateTime(), nullable=False), sa.Column("status", sa.String(length=20), nullable=False), From 275bc0cb535d151488c247bd7f6cfb411b5026e4 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 18:21:57 +0800 Subject: [PATCH 29/51] =?UTF-8?q?:recycles:=20=E6=B3=A8=E9=87=8A=E6=8E=89c?= =?UTF-8?q?ookie=E7=9B=B8=E5=85=B3=E4=BB=A3=E7=A0=81=EF=BC=8C=E4=BD=BF?= =?UTF-8?q?=E5=BE=97bison=E5=8F=AF=E4=BB=A5=E6=AD=A3=E5=B8=B8=E8=BF=90?= =?UTF-8?q?=E8=A1=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 248 +++++++++++++------------- nonebot_bison/sub_manager/__init__.py | 78 ++++---- nonebot_bison/utils/site.py | 219 ++++++++++++----------- 3 files changed, 275 insertions(+), 270 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 5cb0152..4fdea77 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,130 +259,130 @@ class DBConfig: ) return res - async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: - """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" - async with create_session() as sess: - query = select(Cookie).distinct().where(Cookie.is_universal == False) # 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 get_unviersal_cookie(self, platform_name: str = None) -> Sequence[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) - res = (await sess.scalars(query)).all() - 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) - await sess.commit() - 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)) - if not cookie_in_db: - return - cookie_in_db.content = cookie.content - cookie_in_db.last_usage = cookie.last_usage - cookie_in_db.status = cookie.status - cookie_in_db.tags = cookie.tags - await sess.commit() - - async def delete_cookie_by_id(self, cookie_id: int): - async with create_session() as sess: - cookie = await sess.scalar( - select(Cookie) - .where(Cookie.id == cookie_id) - .outerjoin(CookieTarget) - .options(selectinload(Cookie.targets)) - ) - if len(cookie.targets) > 0: - raise Exception(f"cookie {cookie.id} in use") - await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) - await sess.commit() - - async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: - async with create_session() as sess: - query = ( - select(Cookie) - .join(CookieTarget) - .join(Target) - .where(Target.platform_name == platform_name, Target.target == target) - ) - return (await sess.scalars(query)).all() - - async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]: - async with create_session() as sess: - query = ( - select(Cookie) - .where(Cookie.platform_name == platform_name) - .where(Cookie.is_universal == True) # noqa: E712 - ) - return (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( - select(Target).where(Target.platform_name == platform_name, Target.target == target) - ) - # check if relation exists - cookie_target = await sess.scalar( - select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) - ) - if cookie_target: - raise DuplicateCookieTargetException() - cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) - cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) - sess.add(cookie_target) - await sess.commit() - - async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): - async with create_session() as sess: - target_obj = await sess.scalar( - select(Target).where(Target.platform_name == platform_name, Target.target == target) - ) - cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) - await sess.execute( - delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) - ) - await sess.commit() - - async def delete_cookie_target_by_id(self, cookie_target_id: int): - async with create_session() as sess: - await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) - await sess.commit() - - async def get_cookie_target(self) -> list[CookieTarget]: - async with create_session() as sess: - query = ( - select(CookieTarget) - .outerjoin(Target) - .options(selectinload(CookieTarget.target)) - .outerjoin(Cookie) - .options(selectinload(CookieTarget.cookie)) - ) - res = list((await sess.scalars(query)).all()) - res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) - return res + # async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: + # """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" + # async with create_session() as sess: + # query = select(Cookie).distinct().where(Cookie.is_universal == False) # 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 get_unviersal_cookie(self, platform_name: str = None) -> Sequence[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) + # res = (await sess.scalars(query)).all() + # 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) + # await sess.commit() + # 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)) + # if not cookie_in_db: + # return + # cookie_in_db.content = cookie.content + # cookie_in_db.last_usage = cookie.last_usage + # cookie_in_db.status = cookie.status + # cookie_in_db.tags = cookie.tags + # await sess.commit() + # + # async def delete_cookie_by_id(self, cookie_id: int): + # async with create_session() as sess: + # cookie = await sess.scalar( + # select(Cookie) + # .where(Cookie.id == cookie_id) + # .outerjoin(CookieTarget) + # .options(selectinload(Cookie.targets)) + # ) + # if len(cookie.targets) > 0: + # raise Exception(f"cookie {cookie.id} in use") + # await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) + # await sess.commit() + # + # async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: + # async with create_session() as sess: + # query = ( + # select(Cookie) + # .join(CookieTarget) + # .join(Target) + # .where(Target.platform_name == platform_name, Target.target == target) + # ) + # return (await sess.scalars(query)).all() + # + # async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]: + # async with create_session() as sess: + # query = ( + # select(Cookie) + # .where(Cookie.platform_name == platform_name) + # .where(Cookie.is_universal == True) # noqa: E712 + # ) + # return (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( + # select(Target).where(Target.platform_name == platform_name, Target.target == target) + # ) + # # check if relation exists + # cookie_target = await sess.scalar( + # select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) + # ) + # if cookie_target: + # raise DuplicateCookieTargetException() + # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + # cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) + # sess.add(cookie_target) + # await sess.commit() + # + # async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): + # async with create_session() as sess: + # target_obj = await sess.scalar( + # select(Target).where(Target.platform_name == platform_name, Target.target == target) + # ) + # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + # await sess.execute( + # delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) + # ) + # await sess.commit() + # + # async def delete_cookie_target_by_id(self, cookie_target_id: int): + # async with create_session() as sess: + # await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) + # await sess.commit() + # + # async def get_cookie_target(self) -> list[CookieTarget]: + # async with create_session() as sess: + # query = ( + # select(CookieTarget) + # .outerjoin(Target) + # .options(selectinload(CookieTarget.target)) + # .outerjoin(Cookie) + # .options(selectinload(CookieTarget.cookie)) + # ) + # res = list((await sess.scalars(query)).all()) + # res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) + # return res config = DBConfig() diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 943efd1..ec58e30 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -14,10 +14,10 @@ from nonebot.adapters.onebot.v11.event import PrivateMessageEvent from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub -from .add_cookie import do_add_cookie -from .del_cookie import do_del_cookie -from .add_cookie_target import do_add_cookie_target -from .del_cookie_target import do_del_cookie_target +# from .add_cookie import do_add_cookie +# from .del_cookie import do_del_cookie +# from .add_cookie_target import do_add_cookie_target +# from .del_cookie_target import do_del_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info add_sub_matcher = on_command( @@ -44,41 +44,41 @@ del_sub_matcher = on_command( del_sub_matcher.handle()(set_target_user_info) do_del_sub(del_sub_matcher) -add_cookie_matcher = on_command( - "添加cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_add_cookie(add_cookie_matcher) - -add_cookie_target_matcher = on_command( - "关联cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_add_cookie_target(add_cookie_target_matcher) - -del_cookie_target_matcher = on_command( - "取消关联cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_del_cookie_target(del_cookie_target_matcher) - -del_cookie_matcher = on_command( - "删除cookie", - rule=configurable_to_me, - permission=SUPERUSER, - priority=5, - block=True, -) -do_del_cookie(del_cookie_matcher) +# add_cookie_matcher = on_command( +# "添加cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_add_cookie(add_cookie_matcher) +# +# add_cookie_target_matcher = on_command( +# "关联cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_add_cookie_target(add_cookie_target_matcher) +# +# del_cookie_target_matcher = on_command( +# "取消关联cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_del_cookie_target(del_cookie_target_matcher) +# +# del_cookie_matcher = on_command( +# "删除cookie", +# rule=configurable_to_me, +# permission=SUPERUSER, +# priority=5, +# block=True, +# ) +# do_del_cookie(del_cookie_matcher) group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index b53aab2..b93ff56 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -45,117 +45,122 @@ class DefaultClientManager(ClientManager): def is_cookie_client_manager(manger: type[ClientManager]) -> bool: return hasattr(manger, "_cookie_client_manger_") - +# TODO: Implement CookieClientManager class CookieClientManager(ClientManager): - _cookie_client_manger_ = True +# _cookie_client_manger_ = True _platform_name: str - _cookie_cd: int = 10 +# _cookie_cd: int = 10 +# +# @classmethod +# async def init_universal_cookie(cls): +# """移除已有的匿名cookie，添加一个新的匿名cookie""" +# universal_cookies = await config.get_unviersal_cookie(cls._platform_name) +# universal_cookie = Cookie( +# platform_name=cls._platform_name, content="{}", is_universal=True, tags={"temporary": True} +# ) +# for cookie in universal_cookies: +# if not cookie.tags.get("temporary"): +# continue +# await config.delete_cookie_by_id(cookie.id) +# universal_cookie.id = cookie.id # 保持原有的id +# await config.add_cookie(universal_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 +# +# @classmethod +# async def valid_cookie(cls, content: str) -> bool: +# """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" +# try: +# data = json.loads(content) +# if not isinstance(data, dict): +# return False +# except JSONDecodeError: +# return False +# return True +# +# @classmethod +# async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: +# """获取 cookie 的友好名字，用于展示""" +# from . import text_fletten +# +# return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") +# +# def _generate_hook(self, cookie: Cookie) -> callable: +# """hook 函数生成器，用于回写请求状态到数据库""" +# +# async def _response_hook(resp: httpx.Response): +# if resp.status_code == 200: +# logger.debug(f"请求成功: {cookie.id} {resp.request.url}") +# cookie.status = "success" +# else: +# logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") +# cookie.status = "failed" +# cookie.last_usage = datetime.now() +# await config.update_cookie(cookie) +# +# return _response_hook +# +# async def _choose_cookie(self, target: Target) -> Cookie: +# """选择 cookie 的具体算法""" +# cookies = await config.get_universal_cookie(self._platform_name) +# if target: +# cookies += await config.get_cookie(self._platform_name, target) +# 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: +# """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控""" +# return cookie +# +# async def get_client(self, target: Target | None) -> AsyncClient: +# """获取 client，根据 target 选择 cookie""" +# client = http_client() +# cookie = await self._choose_cookie(target) +# if cookie.is_universal: +# logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie") +# else: +# logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") +# +# return await self._assemble_client(client, cookie) +# +# async def _assemble_client(self, client, cookie) -> AsyncClient: +# """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" +# cookies = httpx.Cookies() +# if cookie: +# cookies.update(json.loads(cookie.content)) +# client.cookies = cookies +# client._bison_cookie = cookie +# client.event_hooks = {"response": [self._generate_hook(cookie)]} +# return client +# +# async def get_client_for_static(self) -> AsyncClient: +# return http_client() +# +# async def get_query_name_client(self) -> AsyncClient: +# return http_client() +# +# async def refresh_client(self): +# pass +# +# +# def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: +# """创建一个平台特化的 CookieClientManger""" +# return type( +# "CookieClientManager", +# (CookieClientManager,), +# {"_platform_name": platform_name}, +# ) - @classmethod - async def init_universal_cookie(cls): - """移除已有的匿名cookie，添加一个新的匿名cookie""" - universal_cookies = await config.get_unviersal_cookie(cls._platform_name) - universal_cookie = Cookie( - platform_name=cls._platform_name, content="{}", is_universal=True, tags={"temporary": True} - ) - for cookie in universal_cookies: - if not cookie.tags.get("temporary"): - continue - await config.delete_cookie_by_id(cookie.id) - universal_cookie.id = cookie.id # 保持原有的id - await config.add_cookie(universal_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 - - @classmethod - async def valid_cookie(cls, content: str) -> bool: - """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" - try: - data = json.loads(content) - if not isinstance(data, dict): - return False - except JSONDecodeError: - return False - return True - - @classmethod - async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: - """获取 cookie 的友好名字，用于展示""" - from . import text_fletten - - return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") - - def _generate_hook(self, cookie: Cookie) -> callable: - """hook 函数生成器，用于回写请求状态到数据库""" - - async def _response_hook(resp: httpx.Response): - if resp.status_code == 200: - logger.debug(f"请求成功: {cookie.id} {resp.request.url}") - cookie.status = "success" - else: - logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") - cookie.status = "failed" - cookie.last_usage = datetime.now() - await config.update_cookie(cookie) - - return _response_hook - - async def _choose_cookie(self, target: Target) -> Cookie: - """选择 cookie 的具体算法""" - cookies = await config.get_universal_cookie(self._platform_name) - if target: - cookies += await config.get_cookie(self._platform_name, target) - 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: - """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控""" - return cookie - - async def get_client(self, target: Target | None) -> AsyncClient: - """获取 client，根据 target 选择 cookie""" - client = http_client() - cookie = await self._choose_cookie(target) - if cookie.is_universal: - logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie") - else: - logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") - - return await self._assemble_client(client, cookie) - - async def _assemble_client(self, client, cookie) -> AsyncClient: - """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" - cookies = httpx.Cookies() - if cookie: - cookies.update(json.loads(cookie.content)) - client.cookies = cookies - client._bison_cookie = cookie - client.event_hooks = {"response": [self._generate_hook(cookie)]} - return client - - async def get_client_for_static(self) -> AsyncClient: - return http_client() - - async def get_query_name_client(self) -> AsyncClient: - return http_client() - - async def refresh_client(self): - pass - - -def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: +def create_cookie_client_manager(platform_name: str): """创建一个平台特化的 CookieClientManger""" - return type( - "CookieClientManager", - (CookieClientManager,), - {"_platform_name": platform_name}, - ) + return DefaultClientManager + class Site: From 4db7e7b911e2dcdd7aa506772a6283a81857f374 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 8 Sep 2024 18:38:38 +0800 Subject: [PATCH 30/51] =?UTF-8?q?:recycles:=20DBConfig=E4=B8=AD=20?= =?UTF-8?q?=E6=9B=BF=E6=8D=A2platform=5Fname=E4=B8=BAsite=5Fname?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 248 +++++++++++++++--------------- 1 file changed, 124 insertions(+), 124 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 4fdea77..4bb58ea 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,130 +259,130 @@ class DBConfig: ) return res - # async def get_cookie(self, platform_name: str = None, target: T_Target = None) -> Sequence[Cookie]: - # """根据平台名和订阅名获取 cookie，不会返回匿名cookie""" - # async with create_session() as sess: - # query = select(Cookie).distinct().where(Cookie.is_universal == False) # 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 get_unviersal_cookie(self, platform_name: str = None) -> Sequence[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) - # res = (await sess.scalars(query)).all() - # 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) - # await sess.commit() - # 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)) - # if not cookie_in_db: - # return - # cookie_in_db.content = cookie.content - # cookie_in_db.last_usage = cookie.last_usage - # cookie_in_db.status = cookie.status - # cookie_in_db.tags = cookie.tags - # await sess.commit() - # - # async def delete_cookie_by_id(self, cookie_id: int): - # async with create_session() as sess: - # cookie = await sess.scalar( - # select(Cookie) - # .where(Cookie.id == cookie_id) - # .outerjoin(CookieTarget) - # .options(selectinload(Cookie.targets)) - # ) - # if len(cookie.targets) > 0: - # raise Exception(f"cookie {cookie.id} in use") - # await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) - # await sess.commit() - # - # async def get_cookie_by_target(self, target: T_Target, platform_name: str) -> Sequence[Cookie]: - # async with create_session() as sess: - # query = ( - # select(Cookie) - # .join(CookieTarget) - # .join(Target) - # .where(Target.platform_name == platform_name, Target.target == target) - # ) - # return (await sess.scalars(query)).all() - # - # async def get_universal_cookie(self, platform_name: str) -> Sequence[Cookie]: - # async with create_session() as sess: - # query = ( - # select(Cookie) - # .where(Cookie.platform_name == platform_name) - # .where(Cookie.is_universal == True) # noqa: E712 - # ) - # return (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( - # select(Target).where(Target.platform_name == platform_name, Target.target == target) - # ) - # # check if relation exists - # cookie_target = await sess.scalar( - # select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) - # ) - # if cookie_target: - # raise DuplicateCookieTargetException() - # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) - # cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) - # sess.add(cookie_target) - # await sess.commit() - # - # async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): - # async with create_session() as sess: - # target_obj = await sess.scalar( - # select(Target).where(Target.platform_name == platform_name, Target.target == target) - # ) - # cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) - # await sess.execute( - # delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) - # ) - # await sess.commit() - # - # async def delete_cookie_target_by_id(self, cookie_target_id: int): - # async with create_session() as sess: - # await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) - # await sess.commit() - # - # async def get_cookie_target(self) -> list[CookieTarget]: - # async with create_session() as sess: - # query = ( - # select(CookieTarget) - # .outerjoin(Target) - # .options(selectinload(CookieTarget.target)) - # .outerjoin(Cookie) - # .options(selectinload(CookieTarget.cookie)) - # ) - # res = list((await sess.scalars(query)).all()) - # res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) - # return res + async def get_cookie(self, site_name: str = None, target: T_Target = None) -> Sequence[Cookie]: + """根据 site_name 和 target 获取 cookie，不会返回匿名cookie""" + async with create_session() as sess: + query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 + if site_name: + query = query.where(Cookie.site_name == site_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 get_unviersal_cookie(self, site_name: str = None) -> Sequence[Cookie]: + async with create_session() as sess: + query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 + if site_name: + query = query.where(Cookie.site_name == site_name) + res = (await sess.scalars(query)).all() + return res + + async def add_cookie_with_content(self, site_name: str, content: str) -> int: + async with create_session() as sess: + cookie = Cookie(site_name=site_name, content=content) + sess.add(cookie) + await sess.commit() + 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)) + if not cookie_in_db: + return + cookie_in_db.content = cookie.content + cookie_in_db.last_usage = cookie.last_usage + cookie_in_db.status = cookie.status + cookie_in_db.tags = cookie.tags + await sess.commit() + + async def delete_cookie_by_id(self, cookie_id: int): + async with create_session() as sess: + cookie = await sess.scalar( + select(Cookie) + .where(Cookie.id == cookie_id) + .outerjoin(CookieTarget) + .options(selectinload(Cookie.targets)) + ) + if len(cookie.targets) > 0: + raise Exception(f"cookie {cookie.id} in use") + await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) + await sess.commit() + + async def get_cookie_by_target(self, target: T_Target, site_name: str) -> Sequence[Cookie]: + async with create_session() as sess: + query = ( + select(Cookie) + .join(CookieTarget) + .join(Target) + .where(Target.site_name == site_name, Target.target == target) + ) + return (await sess.scalars(query)).all() + + async def get_universal_cookie(self, site_name: str) -> Sequence[Cookie]: + async with create_session() as sess: + query = ( + select(Cookie) + .where(Cookie.site_name == site_name) + .where(Cookie.is_universal == True) # noqa: E712 + ) + return (await sess.scalars(query)).all() + + async def add_cookie_target(self, target: T_Target, site_name: str, cookie_id: int): + async with create_session() as sess: + target_obj = await sess.scalar( + select(Target).where(Target.site_name == site_name, Target.target == target) + ) + # check if relation exists + cookie_target = await sess.scalar( + select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id) + ) + if cookie_target: + raise DuplicateCookieTargetException() + cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj) + sess.add(cookie_target) + await sess.commit() + + async def delete_cookie_target(self, target: T_Target, site_name: str, cookie_id: int): + async with create_session() as sess: + target_obj = await sess.scalar( + select(Target).where(Target.site_name == site_name, Target.target == target) + ) + cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + await sess.execute( + delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) + ) + await sess.commit() + + async def delete_cookie_target_by_id(self, cookie_target_id: int): + async with create_session() as sess: + await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id)) + await sess.commit() + + async def get_cookie_target(self) -> list[CookieTarget]: + async with create_session() as sess: + query = ( + select(CookieTarget) + .outerjoin(Target) + .options(selectinload(CookieTarget.target)) + .outerjoin(Cookie) + .options(selectinload(CookieTarget.cookie)) + ) + res = list((await sess.scalars(query)).all()) + res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id)) + return res config = DBConfig() From f959e3ee0886eefab283a15548d522d021043c1d Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 10:04:19 +0800 Subject: [PATCH 31/51] =?UTF-8?q?:recycles:=20=E4=BB=BF=E7=85=A7=20platfor?= =?UTF-8?q?m=5Fmanager=20=E6=B7=BB=E5=8A=A0=20site=5Fmanager?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/platform/__init__.py | 9 +++++++++ nonebot_bison/platform/platform.py | 17 +---------------- nonebot_bison/sub_manager/__init__.py | 16 ++++++++-------- nonebot_bison/sub_manager/add_cookie.py | 21 +++++++++++---------- nonebot_bison/types.py | 15 +++++++++++++++ nonebot_bison/utils/__init__.py | 18 +++++++++--------- nonebot_bison/utils/site.py | 7 +++++-- 7 files changed, 58 insertions(+), 45 deletions(-) diff --git a/nonebot_bison/platform/__init__.py b/nonebot_bison/platform/__init__.py index c99ce12..66a6e5b 100644 --- a/nonebot_bison/platform/__init__.py +++ b/nonebot_bison/platform/__init__.py @@ -4,6 +4,7 @@ from collections import defaultdict from importlib import import_module from .platform import Platform, make_no_target_group +from ..utils import Site _package_dir = str(Path(__file__).resolve().parent) for _, module_name, _ in iter_modules([_package_dir]): @@ -22,3 +23,11 @@ for name, platform_list in _platform_list.items(): platform_manager[name] = platform_list[0] else: platform_manager[name] = make_no_target_group(platform_list) + + +site_manager: dict[str, type[Site]] = {} +for site in Site.registry: + if not hasattr(site, "name"): + continue + site_manager[site.name] = site + diff --git a/nonebot_bison/platform/platform.py b/nonebot_bison/platform/platform.py index ecbafe6..111adaf 100644 --- a/nonebot_bison/platform/platform.py +++ b/nonebot_bison/platform/platform.py @@ -16,7 +16,7 @@ from nonebot_plugin_saa import PlatformTarget from ..post import Post from ..utils import Site, ProcessContext from ..plugin_config import plugin_config -from ..types import Tag, Target, RawPost, SubUnit, Category +from ..types import Tag, Target, RawPost, SubUnit, Category, RegistryMeta class CategoryNotSupport(Exception): @@ -29,21 +29,6 @@ class CategoryNotRecognize(Exception): """raise in get_category, when you don't know the category of post""" -class RegistryMeta(type): - def __new__(cls, name, bases, namespace, **kwargs): - return super().__new__(cls, name, bases, namespace) - - def __init__(cls, name, bases, namespace, **kwargs): - if kwargs.get("base"): - # this is the base class - cls.registry = [] - elif not kwargs.get("abstract"): - # this is the subclass - cls.registry.append(cls) - - super().__init__(name, bases, namespace, **kwargs) - - P = ParamSpec("P") R = TypeVar("R") diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index ec58e30..d54c0f8 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -44,15 +44,15 @@ del_sub_matcher = on_command( del_sub_matcher.handle()(set_target_user_info) do_del_sub(del_sub_matcher) -# add_cookie_matcher = on_command( -# "添加cookie", -# rule=configurable_to_me, -# permission=SUPERUSER, -# priority=5, -# block=True, -# ) +add_cookie_matcher = on_command( + "添加cookie", + rule=configurable_to_me, + permission=SUPERUSER, + priority=5, + block=True, +) # do_add_cookie(add_cookie_matcher) -# + # add_cookie_target_matcher = on_command( # "关联cookie", # rule=configurable_to_me, diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 9ae36e3..4b912fe 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -1,14 +1,14 @@ from typing import cast -from nonebot.typing import T_State +from nonebot.adapters import Message, MessageTemplate from nonebot.matcher import Matcher from nonebot.params import Arg, ArgPlainText -from nonebot.adapters import Message, MessageTemplate +from nonebot.typing import T_State +from .utils import common_platform, gen_handle_cancel 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 @@ -20,12 +20,12 @@ def do_add_cookie(add_cookie: type[Matcher]): state["_prompt"] = ( "请输入想要添加 Cookie 的平台，目前支持，请输入冒号左边的名称：\n" + "".join( - [ - f"{platform_name}: {platform_manager[platform_name].name}\n" - for platform_name in common_platform - if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) - ] - ) + [ + f"{platform_name}: {platform_manager[platform_name].name}\n" + for platform_name in common_platform + if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) + ] + ) + "要查看全部平台请输入：“全部”\n中止添加cookie过程请输入：“取消”" ) @@ -53,7 +53,8 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.got("cookie", MessageTemplate("{_prompt}"), [handle_cancel]) async def got_cookie(state: T_State, cookie: Message = Arg()): - client_mgr: CookieClientManager = platform_manager[state["platform"]].site.client_mgr + client_mgr: type[CookieClientManager] = cast(type[CookieClientManager], + platform_manager[state["platform"]].site.client_mgr) cookie_text = cookie.extract_plain_text() if not await client_mgr.valid_cookie(cookie_text): await add_cookie.reject("无效的 Cookie，请检查后重新输入，详情见<待添加的文档>") diff --git a/nonebot_bison/types.py b/nonebot_bison/types.py index 0d08bfd..d90bcd3 100644 --- a/nonebot_bison/types.py +++ b/nonebot_bison/types.py @@ -58,3 +58,18 @@ class ApiError(Exception): class SubUnit(NamedTuple): sub_target: Target user_sub_infos: list[UserSubInfo] + + +class RegistryMeta(type): + def __new__(cls, name, bases, namespace, **kwargs): + return super().__new__(cls, name, bases, namespace) + + def __init__(cls, name, bases, namespace, **kwargs): + if kwargs.get("base"): + # this is the base class + cls.registry = [] + elif not kwargs.get("abstract"): + # this is the subclass + cls.registry.append(cls) + + super().__init__(name, bases, namespace, **kwargs) diff --git a/nonebot_bison/utils/__init__.py b/nonebot_bison/utils/__init__.py index 5da51f3..c61ef46 100644 --- a/nonebot_bison/utils/__init__.py +++ b/nonebot_bison/utils/__init__.py @@ -1,25 +1,25 @@ +import difflib import re import sys -import difflib import nonebot -from nonebot.plugin import require from bs4 import BeautifulSoup as bs from nonebot.log import logger, default_format +from nonebot.plugin import require from nonebot_plugin_saa import Text, Image, MessageSegmentFactory -from .site import Site as Site -from ..plugin_config import plugin_config -from .image import pic_merge as pic_merge +from .context import ProcessContext as ProcessContext from .http import http_client as http_client from .image import capture_html as capture_html -from .site import ClientManager as ClientManager -from .image import text_to_image as text_to_image -from .site import anonymous_site as anonymous_site -from .context import ProcessContext as ProcessContext from .image import is_pics_mergable as is_pics_mergable +from .image import pic_merge as pic_merge from .image import pic_url_to_image as pic_url_to_image +from .image import text_to_image as text_to_image +from .site import ClientManager as ClientManager from .site import DefaultClientManager as DefaultClientManager +from .site import Site as Site +from .site import anonymous_site as anonymous_site +from ..plugin_config import plugin_config class Singleton(type): diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index b93ff56..1ad2f4e 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -7,8 +7,9 @@ from datetime import datetime, timedelta import httpx from httpx import AsyncClient from nonebot.log import logger +from sqlalchemy.testing.suite.test_reflection import metadata -from ..types import Target +from ..types import Target, RegistryMeta from ..config import config from .http import http_client from ..config.db_model import Cookie @@ -163,12 +164,14 @@ def create_cookie_client_manager(platform_name: str): -class Site: +class Site(metaclass=RegistryMeta,base=True): schedule_type: Literal["date", "interval", "cron"] schedule_setting: dict name: str client_mgr: type[ClientManager] = DefaultClientManager require_browser: bool = False + registry: list[type["Site"]] + def __str__(self): return f"[{self.name}]-{self.name}-{self.schedule_setting}" From 65a5976897382876178cebea859bf7f55cc3c6c9 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 11:06:59 +0800 Subject: [PATCH 32/51] :recycles: CookieClientManager --- nonebot_bison/utils/site.py | 225 ++++++++++++++++++------------------ 1 file changed, 112 insertions(+), 113 deletions(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 1ad2f4e..69503b7 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -46,122 +46,121 @@ class DefaultClientManager(ClientManager): def is_cookie_client_manager(manger: type[ClientManager]) -> bool: return hasattr(manger, "_cookie_client_manger_") -# TODO: Implement CookieClientManager class CookieClientManager(ClientManager): -# _cookie_client_manger_ = True - _platform_name: str -# _cookie_cd: int = 10 -# -# @classmethod -# async def init_universal_cookie(cls): -# """移除已有的匿名cookie，添加一个新的匿名cookie""" -# universal_cookies = await config.get_unviersal_cookie(cls._platform_name) -# universal_cookie = Cookie( -# platform_name=cls._platform_name, content="{}", is_universal=True, tags={"temporary": True} -# ) -# for cookie in universal_cookies: -# if not cookie.tags.get("temporary"): -# continue -# await config.delete_cookie_by_id(cookie.id) -# universal_cookie.id = cookie.id # 保持原有的id -# await config.add_cookie(universal_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 -# -# @classmethod -# async def valid_cookie(cls, content: str) -> bool: -# """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" -# try: -# data = json.loads(content) -# if not isinstance(data, dict): -# return False -# except JSONDecodeError: -# return False -# return True -# -# @classmethod -# async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: -# """获取 cookie 的友好名字，用于展示""" -# from . import text_fletten -# -# return text_fletten(f"{cookie.platform_name} [{cookie.content[:10]}]") -# -# def _generate_hook(self, cookie: Cookie) -> callable: -# """hook 函数生成器，用于回写请求状态到数据库""" -# -# async def _response_hook(resp: httpx.Response): -# if resp.status_code == 200: -# logger.debug(f"请求成功: {cookie.id} {resp.request.url}") -# cookie.status = "success" -# else: -# logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") -# cookie.status = "failed" -# cookie.last_usage = datetime.now() -# await config.update_cookie(cookie) -# -# return _response_hook -# -# async def _choose_cookie(self, target: Target) -> Cookie: -# """选择 cookie 的具体算法""" -# cookies = await config.get_universal_cookie(self._platform_name) -# if target: -# cookies += await config.get_cookie(self._platform_name, target) -# 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: -# """检查Cookie，可以做一些自定义的逻辑，比如说Site的统一风控""" -# return cookie -# -# async def get_client(self, target: Target | None) -> AsyncClient: -# """获取 client，根据 target 选择 cookie""" -# client = http_client() -# cookie = await self._choose_cookie(target) -# if cookie.is_universal: -# logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie") -# else: -# logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}") -# -# return await self._assemble_client(client, cookie) -# -# async def _assemble_client(self, client, cookie) -> AsyncClient: -# """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" -# cookies = httpx.Cookies() -# if cookie: -# cookies.update(json.loads(cookie.content)) -# client.cookies = cookies -# client._bison_cookie = cookie -# client.event_hooks = {"response": [self._generate_hook(cookie)]} -# return client -# -# async def get_client_for_static(self) -> AsyncClient: -# return http_client() -# -# async def get_query_name_client(self) -> AsyncClient: -# return http_client() -# -# async def refresh_client(self): -# pass -# -# -# def create_cookie_client_manager(platform_name: str) -> type[CookieClientManager]: -# """创建一个平台特化的 CookieClientManger""" -# return type( -# "CookieClientManager", -# (CookieClientManager,), -# {"_platform_name": platform_name}, -# ) + _cookie_client_manger_ = True + _site_name: str + _cookie_cd: int = 10 -def create_cookie_client_manager(platform_name: str): + @classmethod + async def init_universal_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"): + continue + await config.delete_cookie_by_id(cookie.id) + universal_cookie.id = cookie.id # 保持原有的id + await config.add_cookie(universal_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 + + @classmethod + async def valid_cookie(cls, content: str) -> bool: + """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" + try: + data = json.loads(content) + if not isinstance(data, dict): + return False + except JSONDecodeError: + return False + return True + + @classmethod + async def get_cookie_friendly_name(cls, cookie: Cookie) -> str: + """获取 cookie 的友好名字，用于展示""" + from . import text_fletten + + return text_fletten(f"{cookie.site_name} [{cookie.content[:10]}]") + + def _generate_hook(self, cookie: Cookie) -> callable: + """hook 函数生成器，用于回写请求状态到数据库""" + + async def _response_hook(resp: httpx.Response): + if resp.status_code == 200: + logger.debug(f"请求成功: {cookie.id} {resp.request.url}") + cookie.status = "success" + else: + logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") + cookie.status = "failed" + cookie.last_usage = datetime.now() + await config.update_cookie(cookie) + + return _response_hook + + async def _choose_cookie(self, target: Target) -> Cookie: + """选择 cookie 的具体算法""" + 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()) + 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() + cookie = await self._choose_cookie(target) + if cookie.is_universal: + logger.debug(f"平台 {self._site_name} 未获取到用户cookie, 使用匿名cookie") + else: + logger.debug(f"平台 {self._site_name} 获取到用户cookie: {cookie.id}") + + return await self._assemble_client(client, cookie) + + async def _assemble_client(self, client, cookie) -> AsyncClient: + """组装 client，可以自定义 cookie 对象的 content 装配到 client 中的方式""" + cookies = httpx.Cookies() + if cookie: + cookies.update(json.loads(cookie.content)) + client.cookies = cookies + client._bison_cookie = cookie + client.event_hooks = {"response": [self._generate_hook(cookie)]} + return client + + async def get_client_for_static(self) -> AsyncClient: + return http_client() + + async def get_query_name_client(self) -> AsyncClient: + return http_client() + + async def refresh_client(self): + pass + + +def create_cookie_client_manager(site_name: str) -> type[CookieClientManager]: """创建一个平台特化的 CookieClientManger""" - return DefaultClientManager - + return type( + "CookieClientManager", + (CookieClientManager,), + {"_site_name": site_name}, + ) +# +# def create_cookie_client_manager(platform_name: str): +# """创建一个平台特化的 CookieClientManger""" +# return DefaultClientManager +# class Site(metaclass=RegistryMeta,base=True): From d43d0426188ead159deb7117e31bf1ea4be64a06 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 11:35:39 +0800 Subject: [PATCH 33/51] :recycles: add_cookie --- nonebot_bison/sub_manager/__init__.py | 4 ++-- nonebot_bison/sub_manager/add_cookie.py | 3 ++- nonebot_bison/utils/site.py | 5 ----- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index d54c0f8..59c9294 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -14,7 +14,7 @@ from nonebot.adapters.onebot.v11.event import PrivateMessageEvent from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub -# from .add_cookie import do_add_cookie +from .add_cookie import do_add_cookie # from .del_cookie import do_del_cookie # from .add_cookie_target import do_add_cookie_target # from .del_cookie_target import do_del_cookie_target @@ -51,7 +51,7 @@ add_cookie_matcher = on_command( priority=5, block=True, ) -# do_add_cookie(add_cookie_matcher) +do_add_cookie(add_cookie_matcher) # add_cookie_target_matcher = on_command( # "关联cookie", diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 4b912fe..03dbe5b 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -44,6 +44,7 @@ def do_add_cookie(add_cookie: type[Matcher]): await add_cookie.finish("已中止添加cookie") elif platform in platform_manager: state["platform"] = platform + state["site"] = platform_manager[platform].site else: await add_cookie.reject("平台输入错误") @@ -62,7 +63,7 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.handle() async def add_cookie_process(state: T_State): - cookie = Cookie(platform_name=state["platform"], content=state["cookie"]) + 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) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 69503b7..669b7e6 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -156,11 +156,6 @@ def create_cookie_client_manager(site_name: str) -> type[CookieClientManager]: (CookieClientManager,), {"_site_name": site_name}, ) -# -# def create_cookie_client_manager(platform_name: str): -# """创建一个平台特化的 CookieClientManger""" -# return DefaultClientManager -# class Site(metaclass=RegistryMeta,base=True): From 404b1e445cfd494a690149bfa815aaba731b0b35 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 18:23:23 +0800 Subject: [PATCH 34/51] :recycles: add_cookie_target --- nonebot_bison/config/db_config.py | 19 +++++++++--------- nonebot_bison/platform/weibo.py | 2 +- nonebot_bison/sub_manager/__init__.py | 20 +++++++++---------- .../sub_manager/add_cookie_target.py | 15 +++++++------- nonebot_bison/sub_manager/utils.py | 6 +++--- 5 files changed, 32 insertions(+), 30 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 4bb58ea..163492b 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -1,20 +1,20 @@ import asyncio from collections import defaultdict -from datetime import time, datetime from collections.abc import Callable, Sequence, Awaitable +from datetime import time, datetime from nonebot.compat import model_dump -from sqlalchemy.orm import selectinload -from sqlalchemy.exc import IntegrityError -from sqlalchemy import func, delete, select -from nonebot_plugin_saa import PlatformTarget from nonebot_plugin_datastore import create_session +from nonebot_plugin_saa import PlatformTarget +from sqlalchemy import func, delete, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import selectinload +from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight +from .utils import NoSuchTargetException, DuplicateCookieTargetException +from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp from ..types import Tag from ..types import Target as T_Target -from .utils import NoSuchTargetException, DuplicateCookieTargetException -from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight -from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp def _get_time(): @@ -342,7 +342,8 @@ class DBConfig: async def add_cookie_target(self, target: T_Target, site_name: str, cookie_id: int): async with create_session() as sess: target_obj = await sess.scalar( - select(Target).where(Target.site_name == site_name, Target.target == target) + select(Target).where(Target.target == target) + # TODO: 仅判断 target，可能会有重名现象，还要判断 platform_name ) # check if relation exists cookie_target = await sess.scalar( diff --git a/nonebot_bison/platform/weibo.py b/nonebot_bison/platform/weibo.py index cadcce6..3e3ba6a 100644 --- a/nonebot_bison/platform/weibo.py +++ b/nonebot_bison/platform/weibo.py @@ -40,7 +40,7 @@ class WeiboSite(Site): name = "weibo.com" schedule_type = "interval" schedule_setting = {"seconds": 3} - client_mgr = create_cookie_client_manager("weibo") + client_mgr = create_cookie_client_manager(name) class Weibo(NewMessage): diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 59c9294..a7c5bf6 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -16,7 +16,7 @@ from .del_sub import do_del_sub from .query_sub import do_query_sub from .add_cookie import do_add_cookie # from .del_cookie import do_del_cookie -# from .add_cookie_target import do_add_cookie_target +from .add_cookie_target import do_add_cookie_target # from .del_cookie_target import do_del_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info @@ -53,15 +53,15 @@ add_cookie_matcher = on_command( ) do_add_cookie(add_cookie_matcher) -# add_cookie_target_matcher = on_command( -# "关联cookie", -# rule=configurable_to_me, -# permission=SUPERUSER, -# priority=5, -# block=True, -# ) -# do_add_cookie_target(add_cookie_target_matcher) -# +add_cookie_target_matcher = on_command( + "关联cookie", + rule=configurable_to_me, + permission=SUPERUSER, + priority=5, + block=True, +) +do_add_cookie_target(add_cookie_target_matcher) + # del_cookie_target_matcher = on_command( # "取消关联cookie", # rule=configurable_to_me, diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index 1650428..ab423c3 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -8,7 +8,7 @@ from nonebot.internal.adapter import MessageTemplate from ..config import config from ..utils import parse_text -from ..platform import platform_manager +from ..platform import platform_manager, site_manager from ..utils.site import CookieClientManager from .utils import gen_handle_cancel, generate_sub_list_text @@ -29,16 +29,17 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): try: target_idx = int(target_idx) state["target"] = state["sub_table"][target_idx] + state["site"] = platform_manager[state["target"]["platform_name"]].site except Exception: await add_cookie_target_matcher.reject("序号错误") @add_cookie_target_matcher.handle() async def init_promote_cookie(state: T_State): - cookies = await config.get_cookie(platform_name=state["target"]["platform_name"]) + cookies = await config.get_cookie(site_name=state["site"].name) associated_cookies = await config.get_cookie( target=state["target"]["target"], - platform_name=state["target"]["platform_name"], + site_name=state["site"].name, ) associated_cookie_ids = {cookie.id for cookie in associated_cookies} cookies = [cookie for cookie in cookies if cookie.id not in associated_cookie_ids] @@ -48,7 +49,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): ) state["cookies"] = cookies - client_mgr = cast(CookieClientManager, platform_manager[cookies[0].platform_name].site.client_mgr) + client_mgr = cast(CookieClientManager, state["site"].client_mgr) state["_prompt"] = "请选择一个 Cookie，已关联的 Cookie 不会显示\n" + "\n".join( [f"{idx}. {await client_mgr.get_cookie_friendly_name(cookie)}" for idx, cookie in enumerate(cookies, 1)] ) @@ -63,10 +64,10 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def add_cookie_target_process(state: T_State): - await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id) + await config.add_cookie_target(state["target"]["target"], state["site"].name, state["cookie"].id) cookie = state["cookie"] - client_mgr = cast(CookieClientManager, platform_manager[cookie.platform_name].site.client_mgr) + client_mgr = cast(CookieClientManager, state["site"].client_mgr) await add_cookie_target_matcher.finish( f"已关联 Cookie: {await client_mgr.get_cookie_friendly_name(cookie)} " - f"到订阅 {state['target']['platform_name']} {state['target']['target']}" + f"到订阅 {state['site'].name} {state['target']['target']}" ) diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index 96b080c..ac413d0 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -13,7 +13,7 @@ from nonebot_plugin_saa import PlatformTarget, extract_target from ..config import config from ..types import Category -from ..platform import platform_manager +from ..platform import platform_manager, site_manager from ..plugin_config import plugin_config from ..utils.site import CookieClientManager, is_cookie_client_manager @@ -108,12 +108,12 @@ async def generate_sub_list_text( res += " {}".format(", ".join(sub.tags)) + "\n" if is_show_cookie: target_cookies = await config.get_cookie( - target=sub.target.target, platform_name=sub.target.platform_name + target=sub.target.target, site_name=platform.site.name ) if target_cookies: res += " 关联的 Cookie：\n" for cookie in target_cookies: - client_mgr = cast(CookieClientManager, platform_manager[cookie.platform_name].site.client_mgr) + client_mgr = cast(CookieClientManager, site_manager[platform.site.name].client_mgr) res += f" \t{await client_mgr.get_cookie_friendly_name(cookie)}\n" else: From 4a5e00c0941756d7c91ecbccd9eac7080a178963 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 18:32:53 +0800 Subject: [PATCH 35/51] :recycles: ddel_cookie_target --- nonebot_bison/sub_manager/__init__.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index a7c5bf6..aba8b33 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -17,7 +17,7 @@ from .query_sub import do_query_sub from .add_cookie import do_add_cookie # from .del_cookie import do_del_cookie from .add_cookie_target import do_add_cookie_target -# from .del_cookie_target import do_del_cookie_target +from .del_cookie_target import do_del_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info add_sub_matcher = on_command( @@ -62,15 +62,15 @@ add_cookie_target_matcher = on_command( ) do_add_cookie_target(add_cookie_target_matcher) -# del_cookie_target_matcher = on_command( -# "取消关联cookie", -# rule=configurable_to_me, -# permission=SUPERUSER, -# priority=5, -# block=True, -# ) -# do_del_cookie_target(del_cookie_target_matcher) -# +del_cookie_target_matcher = on_command( + "取消关联cookie", + rule=configurable_to_me, + permission=SUPERUSER, + priority=5, + block=True, +) +do_del_cookie_target(del_cookie_target_matcher) + # del_cookie_matcher = on_command( # "删除cookie", # rule=configurable_to_me, From 2c2c9a091c45be831da38ef310ec9cf62675f44f Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 18:36:18 +0800 Subject: [PATCH 36/51] :recycle: del_cookie --- nonebot_bison/sub_manager/__init__.py | 18 +++++++++--------- nonebot_bison/sub_manager/del_cookie.py | 6 +++--- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index aba8b33..943efd1 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -15,7 +15,7 @@ from .add_sub import do_add_sub from .del_sub import do_del_sub from .query_sub import do_query_sub from .add_cookie import do_add_cookie -# from .del_cookie import do_del_cookie +from .del_cookie import do_del_cookie from .add_cookie_target import do_add_cookie_target from .del_cookie_target import do_del_cookie_target from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info @@ -71,14 +71,14 @@ del_cookie_target_matcher = on_command( ) do_del_cookie_target(del_cookie_target_matcher) -# del_cookie_matcher = on_command( -# "删除cookie", -# rule=configurable_to_me, -# permission=SUPERUSER, -# priority=5, -# block=True, -# ) -# do_del_cookie(del_cookie_matcher) +del_cookie_matcher = on_command( + "删除cookie", + rule=configurable_to_me, + permission=SUPERUSER, + priority=5, + block=True, +) +do_del_cookie(del_cookie_matcher) group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) diff --git a/nonebot_bison/sub_manager/del_cookie.py b/nonebot_bison/sub_manager/del_cookie.py index 4054078..6b3056f 100644 --- a/nonebot_bison/sub_manager/del_cookie.py +++ b/nonebot_bison/sub_manager/del_cookie.py @@ -5,8 +5,8 @@ from nonebot_plugin_saa import MessageFactory from ..config import config from ..utils import parse_text +from ..platform import site_manager from .utils import gen_handle_cancel -from ..platform import platform_manager def do_del_cookie(del_cookie: type[Matcher]): @@ -21,9 +21,9 @@ def do_del_cookie(del_cookie: type[Matcher]): state["cookie_table"] = {} for index, cookie in enumerate(cookies, 1): state["cookie_table"][index] = cookie - client_mgr = platform_manager[cookie.platform_name].site.client_mgr + client_mgr = site_manager[cookie.site_name].client_mgr friendly_name = await client_mgr.get_cookie_friendly_name(cookie) - res += f"{index} {cookie.platform_name} {friendly_name} {len(cookie.targets)}个关联\n" + res += f"{index} {cookie.site_name} {friendly_name} {len(cookie.targets)}个关联\n" if res[-1] != "\n": res += "\n" res += "请输入要删除的 Cookie 的序号\n输入'取消'中止" From f865cef4271152a4bdda7fc25591963d5bf42522 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 12 Sep 2024 23:33:11 +0800 Subject: [PATCH 37/51] =?UTF-8?q?:bug:=20=E8=B0=83=E6=95=B4=E6=97=A5?= =?UTF-8?q?=E5=BF=97=E7=AD=89=E7=BA=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/site.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 669b7e6..ccb1f29 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -7,12 +7,11 @@ from datetime import datetime, timedelta import httpx from httpx import AsyncClient from nonebot.log import logger -from sqlalchemy.testing.suite.test_reflection import metadata -from ..types import Target, RegistryMeta from ..config import config from .http import http_client from ..config.db_model import Cookie +from ..types import Target, RegistryMeta class ClientManager(ABC): @@ -46,6 +45,7 @@ class DefaultClientManager(ClientManager): def is_cookie_client_manager(manger: type[ClientManager]) -> bool: return hasattr(manger, "_cookie_client_manger_") + class CookieClientManager(ClientManager): _cookie_client_manger_ = True _site_name: str @@ -55,9 +55,7 @@ class CookieClientManager(ClientManager): async def init_universal_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} - ) + 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"): continue @@ -95,10 +93,10 @@ class CookieClientManager(ClientManager): async def _response_hook(resp: httpx.Response): if resp.status_code == 200: - logger.debug(f"请求成功: {cookie.id} {resp.request.url}") + logger.trace(f"请求成功: {cookie.id} {resp.request.url}") cookie.status = "success" else: - logger.error(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") + logger.warning(f"请求失败:{cookie.id} {resp.request.url}, 状态码: {resp.status_code}") cookie.status = "failed" cookie.last_usage = datetime.now() await config.update_cookie(cookie) @@ -123,9 +121,9 @@ class CookieClientManager(ClientManager): client = http_client() cookie = await self._choose_cookie(target) if cookie.is_universal: - logger.debug(f"平台 {self._site_name} 未获取到用户cookie, 使用匿名cookie") + logger.trace(f"平台 {self._site_name} 未获取到用户cookie, 使用匿名cookie") else: - logger.debug(f"平台 {self._site_name} 获取到用户cookie: {cookie.id}") + logger.trace(f"平台 {self._site_name} 获取到用户cookie: {cookie.id}") return await self._assemble_client(client, cookie) @@ -158,7 +156,7 @@ def create_cookie_client_manager(site_name: str) -> type[CookieClientManager]: ) -class Site(metaclass=RegistryMeta,base=True): +class Site(metaclass=RegistryMeta, base=True): schedule_type: Literal["date", "interval", "cron"] schedule_setting: dict name: str @@ -166,7 +164,6 @@ class Site(metaclass=RegistryMeta,base=True): require_browser: bool = False registry: list[type["Site"]] - def __str__(self): return f"[{self.name}]-{self.name}-{self.schedule_setting}" From 5111baa89c0a59c9998344824f1cba1880816658 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 00:11:49 +0800 Subject: [PATCH 38/51] =?UTF-8?q?:bug:=20=E5=BA=94=E7=94=A8=E9=83=A8?= =?UTF-8?q?=E5=88=86=E6=8E=A8=E8=8D=90=E7=9A=84=E9=87=8D=E5=91=BD=E5=90=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/scheduler/manager.py | 2 +- nonebot_bison/sub_manager/add_cookie.py | 25 +++++++++++++------------ nonebot_bison/utils/site.py | 5 ++--- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/nonebot_bison/scheduler/manager.py b/nonebot_bison/scheduler/manager.py index 3c4a5a3..852f1bb 100644 --- a/nonebot_bison/scheduler/manager.py +++ b/nonebot_bison/scheduler/manager.py @@ -32,7 +32,7 @@ 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.init_universal_cookie() + await site.client_mgr.refresh_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/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 03dbe5b..b7d35ac 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -1,14 +1,14 @@ from typing import cast -from nonebot.adapters import Message, MessageTemplate +from nonebot.typing import T_State from nonebot.matcher import Matcher from nonebot.params import Arg, ArgPlainText -from nonebot.typing import T_State +from nonebot.adapters import Message, MessageTemplate -from .utils import common_platform, gen_handle_cancel 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 @@ -20,12 +20,12 @@ def do_add_cookie(add_cookie: type[Matcher]): state["_prompt"] = ( "请输入想要添加 Cookie 的平台，目前支持，请输入冒号左边的名称：\n" + "".join( - [ - f"{platform_name}: {platform_manager[platform_name].name}\n" - for platform_name in common_platform - if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) - ] - ) + [ + f"{platform_name}: {platform_manager[platform_name].name}\n" + for platform_name in common_platform + if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) + ] + ) + "要查看全部平台请输入：“全部”\n中止添加cookie过程请输入：“取消”" ) @@ -54,10 +54,11 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.got("cookie", MessageTemplate("{_prompt}"), [handle_cancel]) async def got_cookie(state: T_State, cookie: Message = Arg()): - client_mgr: type[CookieClientManager] = cast(type[CookieClientManager], - platform_manager[state["platform"]].site.client_mgr) + client_mgr: type[CookieClientManager] = cast( + type[CookieClientManager], platform_manager[state["platform"]].site.client_mgr + ) cookie_text = cookie.extract_plain_text() - if not await client_mgr.valid_cookie(cookie_text): + if not await client_mgr.validate_cookie(cookie_text): await add_cookie.reject("无效的 Cookie，请检查后重新输入，详情见<待添加的文档>") state["cookie"] = cookie_text diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index ccb1f29..569b80b 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -52,7 +52,7 @@ class CookieClientManager(ClientManager): _cookie_cd: int = 10 @classmethod - async def init_universal_cookie(cls): + async def refresh_universal_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}) @@ -71,7 +71,7 @@ class CookieClientManager(ClientManager): return cookie @classmethod - async def valid_cookie(cls, content: str) -> bool: + async def validate_cookie(cls, content: str) -> bool: """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写""" try: data = json.loads(content) @@ -133,7 +133,6 @@ class CookieClientManager(ClientManager): if cookie: cookies.update(json.loads(cookie.content)) client.cookies = cookies - client._bison_cookie = cookie client.event_hooks = {"response": [self._generate_hook(cookie)]} return client From 4f73f8a08c98b0598de313448e20511da201a26b Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 00:26:28 +0800 Subject: [PATCH 39/51] =?UTF-8?q?:bug:=20=E5=BC=83=E7=94=A8=5Fcookie=5Fcli?= =?UTF-8?q?ent=5Fmanger=5F=EF=BC=8C=E6=94=B9=E7=94=A8issubclass=E5=88=A4?= =?UTF-8?q?=E6=96=AD=E6=98=AF=E5=90=A6=E4=B8=BACookieClientManager?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/utils/site.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 569b80b..5c1bb8d 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -42,12 +42,7 @@ class DefaultClientManager(ClientManager): pass -def is_cookie_client_manager(manger: type[ClientManager]) -> bool: - return hasattr(manger, "_cookie_client_manger_") - - class CookieClientManager(ClientManager): - _cookie_client_manger_ = True _site_name: str _cookie_cd: int = 10 @@ -146,6 +141,10 @@ class CookieClientManager(ClientManager): pass +def is_cookie_client_manager(manger: type[ClientManager]) -> bool: + return issubclass(manger, CookieClientManager) + + def create_cookie_client_manager(site_name: str) -> type[CookieClientManager]: """创建一个平台特化的 CookieClientManger""" return type( From af246df22201a3246fc680cd60f5089daa85b049 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 00:35:14 +0800 Subject: [PATCH 40/51] =?UTF-8?q?:sparkles:=20=E4=BF=BA=E5=8F=88=E6=9D=A5?= =?UTF-8?q?=E6=94=B9=E6=95=B0=E6=8D=AE=E5=BA=93=E5=93=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_model.py | 14 ++++++++++++-- ...22_add_cookie.py => ef796b74b0fe_add_cookie.py} | 9 +++++---- 2 files changed, 17 insertions(+), 6 deletions(-) rename nonebot_bison/config/migrations/{be215495b122_add_cookie.py => ef796b74b0fe_add_cookie.py} (90%) diff --git a/nonebot_bison/config/db_model.py b/nonebot_bison/config/db_model.py index fcd33bc..54ea825 100644 --- a/nonebot_bison/config/db_model.py +++ b/nonebot_bison/config/db_model.py @@ -79,14 +79,24 @@ class Cookie(Model): # Cookie 当前的状态 status: Mapped[str] = mapped_column(String(20), default="") # 使用一次之后，需要的冷却时间 - cd: Mapped[int] = mapped_column(default=0) - # 是否是通用 Cookie，默认用于匿名 Cookie + cd_milliseconds: Mapped[int] = mapped_column(default=0) + # 是否是通用 Cookie（对所有Target都有效） is_universal: Mapped[bool] = mapped_column(default=False) + # 是否是匿名 Cookie + is_anonymous: 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") + @property + def cd(self) -> datetime.timedelta: + return datetime.timedelta(milliseconds=self.cd_milliseconds) + + @cd.setter + def cd(self, value: datetime.timedelta): + self.cd_milliseconds = int(value.total_seconds() * 1000) + class CookieTarget(Model): id: Mapped[int] = mapped_column(primary_key=True) diff --git a/nonebot_bison/config/migrations/be215495b122_add_cookie.py b/nonebot_bison/config/migrations/ef796b74b0fe_add_cookie.py similarity index 90% rename from nonebot_bison/config/migrations/be215495b122_add_cookie.py rename to nonebot_bison/config/migrations/ef796b74b0fe_add_cookie.py index 87e7bb1..01d5d4a 100644 --- a/nonebot_bison/config/migrations/be215495b122_add_cookie.py +++ b/nonebot_bison/config/migrations/ef796b74b0fe_add_cookie.py @@ -1,8 +1,8 @@ """empty message -Revision ID: be215495b122 +Revision ID: ef796b74b0fe Revises: f9baef347cc8 -Create Date: 2024-09-08 18:12:43.540818 +Create Date: 2024-09-13 00:34:08.601438 """ @@ -12,7 +12,7 @@ from sqlalchemy import Text from sqlalchemy.dialects import postgresql # revision identifiers, used by Alembic. -revision = "be215495b122" +revision = "ef796b74b0fe" down_revision = "f9baef347cc8" branch_labels = None depends_on = None @@ -27,8 +27,9 @@ def upgrade() -> None: sa.Column("content", sa.String(length=1024), nullable=False), sa.Column("last_usage", sa.DateTime(), nullable=False), sa.Column("status", sa.String(length=20), nullable=False), - sa.Column("cd", sa.Integer(), nullable=False), + sa.Column("cd_milliseconds", sa.Integer(), nullable=False), sa.Column("is_universal", sa.Boolean(), nullable=False), + sa.Column("is_anonymous", sa.Boolean(), nullable=False), sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False), sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")), ) From 16331b50d5195cb98c7202a9ced0b19d2aaf335b Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 00:59:19 +0800 Subject: [PATCH 41/51] =?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() From b25fcd9ac292f58344e1b22d3d3e52694f990f00 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 01:07:34 +0800 Subject: [PATCH 42/51] =?UTF-8?q?:sparkles:=20=E5=B0=9D=E8=AF=95=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0=E4=B8=80=E7=A7=8D=E5=8F=AF=E4=BB=A5=E8=B7=B3=E8=BF=87?= =?UTF-8?q?=E5=BD=93=E5=89=8D=E8=AF=B7=E6=B1=82=E7=9A=84=E6=96=B9=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/scheduler/scheduler.py | 3 +++ nonebot_bison/utils/site.py | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/nonebot_bison/scheduler/scheduler.py b/nonebot_bison/scheduler/scheduler.py index cf36a31..f94ea87 100644 --- a/nonebot_bison/scheduler/scheduler.py +++ b/nonebot_bison/scheduler/scheduler.py @@ -12,6 +12,7 @@ from ..send import send_msgs from ..types import Target, SubUnit from ..platform import platform_manager from ..utils import Site, ProcessContext +from ..utils.site import SkipRequestException @dataclass @@ -107,6 +108,8 @@ class Scheduler: schedulable.platform_name, schedulable.target ) to_send = await platform_obj.do_fetch_new_post(SubUnit(schedulable.target, send_userinfo_list)) + except SkipRequestException as err: + logger.debug(f"skip request: {err}") except Exception as err: records = context.gen_req_records() for record in records: diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index e39d1e9..4519f7c 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -173,3 +173,7 @@ def anonymous_site(schedule_type: Literal["date", "interval", "cron"], schedule_ "client_mgr": DefaultClientManager, }, ) + + +class SkipRequestException(Exception): + pass From 4b8d6a93798a01f5eff0f7826042d8c9abd02a3b Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 10:11:54 +0800 Subject: [PATCH 43/51] :bug: fix --- nonebot_bison/config/db_config.py | 32 ++++++++----------- nonebot_bison/sub_manager/__init__.py | 5 ++- .../sub_manager/add_cookie_target.py | 4 +-- nonebot_bison/sub_manager/utils.py | 7 ++-- nonebot_bison/utils/site.py | 2 +- 5 files changed, 23 insertions(+), 27 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 163492b..92da96c 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -1,20 +1,20 @@ import asyncio from collections import defaultdict -from collections.abc import Callable, Sequence, Awaitable from datetime import time, datetime +from collections.abc import Callable, Sequence, Awaitable from nonebot.compat import model_dump -from nonebot_plugin_datastore import create_session -from nonebot_plugin_saa import PlatformTarget -from sqlalchemy import func, delete, select -from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import selectinload +from sqlalchemy.exc import IntegrityError +from sqlalchemy import func, delete, select +from nonebot_plugin_saa import PlatformTarget +from nonebot_plugin_datastore import create_session -from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight -from .utils import NoSuchTargetException, DuplicateCookieTargetException -from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp from ..types import Tag from ..types import Target as T_Target +from .utils import NoSuchTargetException, DuplicateCookieTargetException +from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight +from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp def _get_time(): @@ -332,18 +332,14 @@ class DBConfig: async def get_universal_cookie(self, site_name: str) -> Sequence[Cookie]: async with create_session() as sess: - query = ( - select(Cookie) - .where(Cookie.site_name == site_name) - .where(Cookie.is_universal == True) # noqa: E712 - ) + query = select(Cookie).where(Cookie.site_name == site_name).where(Cookie.is_universal == True) # noqa: E712 return (await sess.scalars(query)).all() - async def add_cookie_target(self, target: T_Target, site_name: str, cookie_id: int): + async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): + """通过 cookie_id 可以唯一确定一个 Cookie，通过 target 和 platform_name 可以唯一确定一个 Target""" async with create_session() as sess: target_obj = await sess.scalar( - select(Target).where(Target.target == target) - # TODO: 仅判断 target，可能会有重名现象，还要判断 platform_name + select(Target).where(Target.platform_name == platform_name, Target.target == target) ) # check if relation exists cookie_target = await sess.scalar( @@ -358,9 +354,7 @@ class DBConfig: async def delete_cookie_target(self, target: T_Target, site_name: str, cookie_id: int): async with create_session() as sess: - target_obj = await sess.scalar( - select(Target).where(Target.site_name == site_name, Target.target == target) - ) + target_obj = await sess.scalar(select(Target).where(Target.site_name == site_name, Target.target == target)) cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) await sess.execute( delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 125afc6..14a924b 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -46,6 +46,7 @@ do_del_sub(del_sub_matcher) add_cookie_matcher = on_command( "添加cookie", + aliases={"添加Cookie"}, rule=configurable_to_me, permission=SUPERUSER, priority=5, @@ -55,6 +56,7 @@ do_add_cookie(add_cookie_matcher) add_cookie_target_matcher = on_command( "关联cookie", + aliases={"关联Cookie"}, rule=configurable_to_me, permission=SUPERUSER, priority=5, @@ -64,6 +66,7 @@ do_add_cookie_target(add_cookie_target_matcher) del_cookie_target_matcher = on_command( "取消关联cookie", + aliases={"取消关联Cookie"}, rule=configurable_to_me, permission=SUPERUSER, priority=5, @@ -73,6 +76,7 @@ do_del_cookie_target(del_cookie_target_matcher) del_cookie_matcher = on_command( "删除cookie", + aliases={"删除Cookie"}, rule=configurable_to_me, permission=SUPERUSER, priority=5, @@ -80,7 +84,6 @@ del_cookie_matcher = on_command( ) do_del_cookie(del_cookie_matcher) - group_manage_matcher = on_command("群管理", rule=to_me(), permission=SUPERUSER, priority=4, block=True) group_handle_cancel = gen_handle_cancel(group_manage_matcher, "已取消") diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index ab423c3..9d31593 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -8,7 +8,7 @@ from nonebot.internal.adapter import MessageTemplate from ..config import config from ..utils import parse_text -from ..platform import platform_manager, site_manager +from ..platform import platform_manager from ..utils.site import CookieClientManager from .utils import gen_handle_cancel, generate_sub_list_text @@ -64,7 +64,7 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def add_cookie_target_process(state: T_State): - await config.add_cookie_target(state["target"]["target"], state["site"].name, state["cookie"].id) + await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id) cookie = state["cookie"] client_mgr = cast(CookieClientManager, state["site"].client_mgr) await add_cookie_target_matcher.finish( diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index ac413d0..ad6cb5e 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -13,8 +13,8 @@ from nonebot_plugin_saa import PlatformTarget, extract_target from ..config import config from ..types import Category -from ..platform import platform_manager, site_manager from ..plugin_config import plugin_config +from ..platform import site_manager, platform_manager from ..utils.site import CookieClientManager, is_cookie_client_manager @@ -75,6 +75,7 @@ async def generate_sub_list_text( is_show_cookie=False, is_hide_no_cookie_platfrom=False, ): + """根据配置参数，生产订阅列表文本，同时将订阅信息存入state["sub_table"]""" if user_info: sub_list = await config.list_subscribe(user_info) else: @@ -107,9 +108,7 @@ async def generate_sub_list_text( if sub.tags: res += " {}".format(", ".join(sub.tags)) + "\n" if is_show_cookie: - target_cookies = await config.get_cookie( - target=sub.target.target, site_name=platform.site.name - ) + target_cookies = await config.get_cookie(target=sub.target.target, site_name=platform.site.name) if target_cookies: res += " 关联的 Cookie：\n" for cookie in target_cookies: diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 4519f7c..2da4422 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -64,7 +64,7 @@ class CookieClientManager(ClientManager): """添加用户 cookie""" cookie = Cookie(site_name=cls._site_name, content=content) cookie.cd = cls._default_cd - config.add_cookie(cookie) + await config.add_cookie(cookie) @classmethod async def validate_cookie(cls, content: str) -> bool: From 4791fb69e0b8f0883c2eaf912455bd241ef444be Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 11:34:45 +0800 Subject: [PATCH 44/51] =?UTF-8?q?:recycle:=20=E9=87=8D=E6=9E=84=20get=5Fco?= =?UTF-8?q?okie=20=E6=96=B9=E6=B3=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/config/db_config.py | 51 +++---- .../sub_manager/add_cookie_target.py | 4 +- nonebot_bison/sub_manager/del_cookie.py | 2 +- nonebot_bison/sub_manager/utils.py | 4 +- nonebot_bison/utils/site.py | 8 +- tests/config/test_cookie.py | 127 +++++++----------- 6 files changed, 73 insertions(+), 123 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 92da96c..2f85775 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -259,36 +259,32 @@ class DBConfig: ) return res - async def get_cookie(self, site_name: str = None, target: T_Target = None) -> Sequence[Cookie]: - """根据 site_name 和 target 获取 cookie，不会返回匿名cookie""" + async def get_cookie( + self, + site_name: str | None = None, + target: T_Target | None = None, + is_universal: bool | None = None, + is_anonymous: bool | None = None, + ) -> Sequence[Cookie]: + """获取满足传入条件的所有 cookie""" async with create_session() as sess: - query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712 + query = select(Cookie).distinct() + if is_universal is not None: + query = query.where(Cookie.is_universal == is_universal) + if is_anonymous is not None: + query = query.where(Cookie.is_anonymous == is_anonymous) if site_name: query = query.where(Cookie.site_name == site_name) query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets)) res = (await sess.scalars(query)).all() if target: + # 如果指定了 target，过滤掉不满足要求的cookie 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] + # 如果指定了 target 且未指定 is_universal，则添加返回 universal cookie + res = [cookie for cookie in res if cookie.id in ids or cookie.is_universal] return res - async def get_unviersal_cookie(self, site_name: str = None) -> Sequence[Cookie]: - async with create_session() as sess: - query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712 - if site_name: - query = query.where(Cookie.site_name == site_name) - res = (await sess.scalars(query)).all() - return res - - async def add_cookie_with_content(self, site_name: str, content: str) -> int: - async with create_session() as sess: - cookie = Cookie(site_name=site_name, content=content) - sess.add(cookie) - await sess.commit() - await sess.refresh(cookie) - return cookie.id - async def add_cookie(self, cookie: Cookie) -> int: async with create_session() as sess: sess.add(cookie) @@ -320,21 +316,6 @@ class DBConfig: await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) await sess.commit() - async def get_cookie_by_target(self, target: T_Target, site_name: str) -> Sequence[Cookie]: - async with create_session() as sess: - query = ( - select(Cookie) - .join(CookieTarget) - .join(Target) - .where(Target.site_name == site_name, Target.target == target) - ) - return (await sess.scalars(query)).all() - - async def get_universal_cookie(self, site_name: str) -> Sequence[Cookie]: - async with create_session() as sess: - query = select(Cookie).where(Cookie.site_name == site_name).where(Cookie.is_universal == True) # noqa: E712 - return (await sess.scalars(query)).all() - async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): """通过 cookie_id 可以唯一确定一个 Cookie，通过 target 和 platform_name 可以唯一确定一个 Target""" async with create_session() as sess: diff --git a/nonebot_bison/sub_manager/add_cookie_target.py b/nonebot_bison/sub_manager/add_cookie_target.py index 9d31593..73d72e2 100644 --- a/nonebot_bison/sub_manager/add_cookie_target.py +++ b/nonebot_bison/sub_manager/add_cookie_target.py @@ -36,10 +36,12 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]): @add_cookie_target_matcher.handle() async def init_promote_cookie(state: T_State): - cookies = await config.get_cookie(site_name=state["site"].name) + # 获取 site 的所有用户 cookie，再排除掉已经关联的 cookie，剩下的就是可以关联的 cookie + cookies = await config.get_cookie(site_name=state["site"].name, is_anonymous=False) associated_cookies = await config.get_cookie( target=state["target"]["target"], site_name=state["site"].name, + is_anonymous=False, ) associated_cookie_ids = {cookie.id for cookie in associated_cookies} cookies = [cookie for cookie in cookies if cookie.id not in associated_cookie_ids] diff --git a/nonebot_bison/sub_manager/del_cookie.py b/nonebot_bison/sub_manager/del_cookie.py index 6b3056f..e343d4f 100644 --- a/nonebot_bison/sub_manager/del_cookie.py +++ b/nonebot_bison/sub_manager/del_cookie.py @@ -14,7 +14,7 @@ def do_del_cookie(del_cookie: type[Matcher]): @del_cookie.handle() async def send_list(state: T_State): - cookies = await config.get_cookie() + cookies = await config.get_cookie(is_anonymous=False) if not cookies: await del_cookie.finish("暂无已添加 Cookie\n请使用“添加cookie”命令添加") res = "已添加的 Cookie 为：\n" diff --git a/nonebot_bison/sub_manager/utils.py b/nonebot_bison/sub_manager/utils.py index ad6cb5e..a74d8ac 100644 --- a/nonebot_bison/sub_manager/utils.py +++ b/nonebot_bison/sub_manager/utils.py @@ -108,7 +108,9 @@ async def generate_sub_list_text( if sub.tags: res += " {}".format(", ".join(sub.tags)) + "\n" if is_show_cookie: - target_cookies = await config.get_cookie(target=sub.target.target, site_name=platform.site.name) + target_cookies = await config.get_cookie( + target=sub.target.target, site_name=platform.site.name, is_anonymous=False + ) if target_cookies: res += " 关联的 Cookie：\n" for cookie in target_cookies: diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 2da4422..b934cf4 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -49,7 +49,7 @@ class CookieClientManager(ClientManager): @classmethod async def refresh_anonymous_cookie(cls): """移除已有的匿名cookie，添加一个新的匿名cookie""" - anonymous_cookies = await config.get_unviersal_cookie(cls._site_name) + anonymous_cookies = await config.get_cookie(cls._site_name, is_anonymous=True) 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: @@ -99,11 +99,9 @@ class CookieClientManager(ClientManager): return _response_hook - async def _choose_cookie(self, target: Target) -> Cookie: + async def _choose_cookie(self, target: Target | None) -> Cookie: """选择 cookie 的具体算法""" - cookies = await config.get_universal_cookie(self._site_name) - if target: - cookies += await config.get_cookie(self._site_name, target) + cookies = await config.get_cookie(self._site_name, target) 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 diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py index b13fa3f..7db0b47 100644 --- a/tests/config/test_cookie.py +++ b/tests/config/test_cookie.py @@ -1,107 +1,74 @@ -import datetime +import json +from typing import cast from nonebug import App -async def test_get_platform_target(app: App, init_scheduler): +async def test_cookie(app: App, init_scheduler): from nonebot_plugin_saa import TargetQQGroup + from nonebot_bison.platform import site_manager from nonebot_bison.config.db_config import config from nonebot_bison.types import Target as T_Target + from nonebot_bison.utils.site import CookieClientManager + target = T_Target("weibo_id") + platform_name = "weibo" await config.add_subscribe( TargetQQGroup(group_id=123), - target=T_Target("weibo_id"), + target=target, target_name="weibo_name", - platform_name="weibo", + platform_name=platform_name, cats=[], tags=[], ) - # await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie") - # cookies = await config.get_cookie_by_user(TargetQQGroup(group_id=123)) - # - # res = await config.get_platform_target("weibo") - # assert len(res) == 2 - # await config.del_subscribe(TargetQQGroup(group_id=123), T_Target("weibo_id1"), "weibo") - # res = await config.get_platform_target("weibo") - # assert len(res) == 2 - # await config.del_subscribe(TargetQQGroup(group_id=123), T_Target("weibo_id"), "weibo") - # res = await config.get_platform_target("weibo") - # assert len(res) == 1 - # - # async with AsyncSession(get_engine()) as sess: - # res = await sess.scalars(select(Target).where(Target.platform_name == "weibo")) - # assert len(res.all()) == 2 - # await config.get_cookie_by_user(TargetQQGroup(group_id=123)) + site = site_manager["weibo.com"] + client_mgr = cast(CookieClientManager, site.client_mgr) + await client_mgr.refresh_anonymous_cookie() # 刷新匿名cookie -async def test_cookie_by_user(app: App, init_scheduler): - from nonebot_plugin_saa import TargetQQGroup + cookies = await config.get_cookie(site_name=site.name) + assert len(cookies) == 1 - from nonebot_bison.config.db_config import config - from nonebot_bison.types import Target as T_Target + await client_mgr.add_user_cookie(json.dumps({"test_cookie": "1"})) + await client_mgr.add_user_cookie(json.dumps({"test_cookie": "2"})) + cookies = await config.get_cookie(site_name=site.name) + assert len(cookies) == 3 + + cookies = await config.get_cookie(site_name=site.name, is_anonymous=False) + assert len(cookies) == 2 + + await config.add_cookie_target(target, platform_name, cookies[0].id) + await config.add_cookie_target(target, platform_name, cookies[1].id) + + cookies = await config.get_cookie(site_name=site.name, target=target) + assert len(cookies) == 3 + + cookies = await config.get_cookie(site_name=site.name, target=target, is_anonymous=False) + assert len(cookies) == 2 + + cookies = await config.get_cookie(site_name=site.name, target=target, is_universal=False) + assert len(cookies) == 2 + + # 测试不同的target + target2 = T_Target("weibo_id2") await config.add_subscribe( TargetQQGroup(group_id=123), - target=T_Target("weibo_id"), - target_name="weibo_name", - platform_name="weibo", + target=target2, + target_name="weibo_name2", + platform_name=platform_name, cats=[], tags=[], ) + await client_mgr.add_user_cookie(json.dumps({"test_cookie": "3"})) + cookies = await config.get_cookie(site_name=site.name, is_anonymous=False) - await config.add_cookie_with_content(TargetQQGroup(group_id=123), "weibo", "cookie") + await config.add_cookie_target(target2, platform_name, cookies[0].id) + await config.add_cookie_target(target2, platform_name, cookies[2].id) - cookies = await config.get_cookie(TargetQQGroup(group_id=123)) - cookie = cookies[0] - assert len(cookies) == 1 - assert cookie.content == "cookie" - assert cookie.platform_name == "weibo" - cookie.last_usage = 0 - assert cookie.status == "" - assert cookie.tags == {} - cookie.content = "cookie1" - cookie.last_usage = datetime.datetime(2024, 8, 22, 0, 0, 0) - cookie.status = "status1" - cookie.tags = {"tag1": "value1"} - await config.update_cookie(cookie) - cookies = await config.get_cookie(TargetQQGroup(group_id=123)) + cookies = await config.get_cookie(site_name=site.name, target=target2) + assert len(cookies) == 3 - assert len(cookies) == 1 - assert cookies[0].content == cookie.content - assert cookies[0].last_usage == cookie.last_usage - assert cookies[0].status == cookie.status - assert cookies[0].tags == cookie.tags - - await config.delete_cookie_by_id(cookies[0].id) - cookies = await config.get_cookie(TargetQQGroup(group_id=123)) - assert len(cookies) == 0 - - -async def test_cookie_target_by_target(app: App, init_scheduler): - from nonebot_plugin_saa import TargetQQGroup - - from nonebot_bison.config.db_config import config - from nonebot_bison.types import Target as T_Target - - await config.add_subscribe( - TargetQQGroup(group_id=123), - target=T_Target("weibo_id"), - target_name="weibo_name", - platform_name="weibo", - cats=[], - tags=[], - ) - - id = await config.add_cookie_with_content(TargetQQGroup(group_id=123), "weibo", "cookie") - - await config.add_cookie_target(T_Target("weibo_id"), "weibo", id) - - cookies = await config.get_cookie_by_target(T_Target("weibo_id"), "weibo") - assert len(cookies) == 1 - assert cookies[0].content == "cookie" - assert cookies[0].platform_name == "weibo" - - await config.delete_cookie_target(T_Target("weibo_id"), "weibo", id) - cookies = await config.get_cookie_by_target(T_Target("weibo_id"), "weibo") - assert len(cookies) == 0 + cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False) + assert len(cookies) == 2 From 2cfd58373f5e4ef99b58bab389a85371da81eb19 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 14:12:02 +0800 Subject: [PATCH 45/51] :white_check_mark: pytest db_config --- nonebot_bison/config/db_config.py | 8 +++-- tests/config/test_cookie.py | 51 ++++++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 2f85775..9dac6dd 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -296,7 +296,7 @@ class DBConfig: async with create_session() as sess: cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id)) if not cookie_in_db: - return + raise ValueError(f"cookie {cookie.id} not found") cookie_in_db.content = cookie.content cookie_in_db.last_usage = cookie.last_usage cookie_in_db.status = cookie.status @@ -333,9 +333,11 @@ class DBConfig: sess.add(cookie_target) await sess.commit() - async def delete_cookie_target(self, target: T_Target, site_name: str, cookie_id: int): + async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int): async with create_session() as sess: - target_obj = await sess.scalar(select(Target).where(Target.site_name == site_name, Target.target == target)) + target_obj = await sess.scalar( + select(Target).where(Target.platform_name == platform_name, Target.target == target) + ) cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) await sess.execute( delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj) diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py index 7db0b47..c831176 100644 --- a/tests/config/test_cookie.py +++ b/tests/config/test_cookie.py @@ -1,6 +1,8 @@ import json from typing import cast +from datetime import datetime +import pytest from nonebug import App @@ -11,6 +13,7 @@ async def test_cookie(app: App, init_scheduler): from nonebot_bison.config.db_config import config from nonebot_bison.types import Target as T_Target from nonebot_bison.utils.site import CookieClientManager + from nonebot_bison.config.utils import DuplicateCookieTargetException target = T_Target("weibo_id") platform_name = "weibo" @@ -25,11 +28,13 @@ async def test_cookie(app: App, init_scheduler): site = site_manager["weibo.com"] client_mgr = cast(CookieClientManager, site.client_mgr) - await client_mgr.refresh_anonymous_cookie() # 刷新匿名cookie + # 刷新匿名cookie + await client_mgr.refresh_anonymous_cookie() cookies = await config.get_cookie(site_name=site.name) assert len(cookies) == 1 + # 添加用户cookie await client_mgr.add_user_cookie(json.dumps({"test_cookie": "1"})) await client_mgr.add_user_cookie(json.dumps({"test_cookie": "2"})) @@ -39,6 +44,7 @@ async def test_cookie(app: App, init_scheduler): cookies = await config.get_cookie(site_name=site.name, is_anonymous=False) assert len(cookies) == 2 + # 单个target，多个cookie await config.add_cookie_target(target, platform_name, cookies[0].id) await config.add_cookie_target(target, platform_name, cookies[1].id) @@ -64,11 +70,54 @@ async def test_cookie(app: App, init_scheduler): await client_mgr.add_user_cookie(json.dumps({"test_cookie": "3"})) cookies = await config.get_cookie(site_name=site.name, is_anonymous=False) + # 多个target，多个cookie await config.add_cookie_target(target2, platform_name, cookies[0].id) await config.add_cookie_target(target2, platform_name, cookies[2].id) cookies = await config.get_cookie(site_name=site.name, target=target2) assert len(cookies) == 3 + # 重复关联 target + with pytest.raises(DuplicateCookieTargetException) as e: + await config.add_cookie_target(target2, platform_name, cookies[2].id) + assert isinstance(e.value, DuplicateCookieTargetException) + cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False) assert len(cookies) == 2 + + # 有关联的cookie不能删除 + with pytest.raises(Exception, match="cookie") as e: + await config.delete_cookie_by_id(cookies[1].id) + cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False) + assert len(cookies) == 2 + + await config.delete_cookie_target(target2, platform_name, cookies[1].id) + await config.delete_cookie_by_id(cookies[1].id) + cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False) + assert len(cookies) == 1 + + cookie = cookies[0] + cookie_id = cookie.id + cookie.last_usage = datetime(2024, 9, 13) + cookie.status = "test" + await config.update_cookie(cookie) + cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False) + assert len(cookies) == 1 + assert cookies[0].id == cookie_id + assert cookies[0].last_usage == datetime(2024, 9, 13) + assert cookies[0].status == "test" + + # 不存在的 cookie_id + cookie.id = 114514 + with pytest.raises(ValueError, match="cookie") as e: + await config.update_cookie(cookie) + + # 获取所有关联对象 + cookie_targets = await config.get_cookie_target() + assert len(cookie_targets) == 3 + + # 删除关联对象 + await config.delete_cookie_target_by_id(cookie_targets[0].id) + + cookie_targets = await config.get_cookie_target() + assert len(cookie_targets) == 2 From dd802a9c17f5f60d329297926cf11f4fa6cfd9f0 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 18:51:28 +0800 Subject: [PATCH 46/51] =?UTF-8?q?:white=5Fcheck=5Fmark:=20=E5=88=9D?= =?UTF-8?q?=E6=AD=A5=E5=AE=9E=E7=8E=B0=E6=B7=BB=E5=8A=A0cookie=E7=9A=84?= =?UTF-8?q?=E5=8D=95=E5=85=83=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/sub_manager/__init__.py | 12 +- nonebot_bison/sub_manager/add_cookie.py | 2 +- nonebot_bison/utils/site.py | 1 + tests/sub_manager/test_add_cookie.py | 212 ++++++++++++++++++++++++ tests/utils.py | 31 ++++ 5 files changed, 253 insertions(+), 5 deletions(-) create mode 100644 tests/sub_manager/test_add_cookie.py diff --git a/nonebot_bison/sub_manager/__init__.py b/nonebot_bison/sub_manager/__init__.py index 14a924b..39dc323 100644 --- a/nonebot_bison/sub_manager/__init__.py +++ b/nonebot_bison/sub_manager/__init__.py @@ -47,7 +47,7 @@ do_del_sub(del_sub_matcher) add_cookie_matcher = on_command( "添加cookie", aliases={"添加Cookie"}, - rule=configurable_to_me, + rule=to_me(), permission=SUPERUSER, priority=5, block=True, @@ -57,7 +57,7 @@ do_add_cookie(add_cookie_matcher) add_cookie_target_matcher = on_command( "关联cookie", aliases={"关联Cookie"}, - rule=configurable_to_me, + rule=to_me(), permission=SUPERUSER, priority=5, block=True, @@ -67,7 +67,7 @@ do_add_cookie_target(add_cookie_target_matcher) del_cookie_target_matcher = on_command( "取消关联cookie", aliases={"取消关联Cookie"}, - rule=configurable_to_me, + rule=to_me(), permission=SUPERUSER, priority=5, block=True, @@ -77,7 +77,7 @@ do_del_cookie_target(del_cookie_target_matcher) del_cookie_matcher = on_command( "删除cookie", aliases={"删除Cookie"}, - rule=configurable_to_me, + rule=to_me(), permission=SUPERUSER, priority=5, block=True, @@ -167,4 +167,8 @@ __all__ = [ "del_sub_matcher", "group_manage_matcher", "no_permission_matcher", + "add_cookie_matcher", + "add_cookie_target_matcher", + "del_cookie_target_matcher", + "del_cookie_matcher", ] diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 20aa004..235c65b 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -57,7 +57,7 @@ def do_add_cookie(add_cookie: type[Matcher]): ) cookie_text = cookie.extract_plain_text() if not await client_mgr.validate_cookie(cookie_text): - await add_cookie.reject("无效的 Cookie，请检查后重新输入，详情见<待添加的文档>") + await add_cookie.reject(state["site"].cookie_format_prompt) state["cookie"] = cookie_text @add_cookie.handle() diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index b934cf4..54716c9 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -156,6 +156,7 @@ class Site(metaclass=RegistryMeta, base=True): client_mgr: type[ClientManager] = DefaultClientManager require_browser: bool = False registry: list[type["Site"]] + cookie_format_prompt = "无效的 Cookie，请检查后重新输入，详情见<待添加的文档>" def __str__(self): return f"[{self.name}]-{self.name}-{self.schedule_setting}" diff --git a/tests/sub_manager/test_add_cookie.py b/tests/sub_manager/test_add_cookie.py new file mode 100644 index 0000000..28c2317 --- /dev/null +++ b/tests/sub_manager/test_add_cookie.py @@ -0,0 +1,212 @@ +import json + +from nonebug.app import App +from pytest_mock import MockerFixture + +from ..utils import BotReply, fake_superuser, fake_admin_user, fake_private_message_event + + +async def test_add_cookie_rule(app: App, mocker: MockerFixture): + from nonebot.adapters.onebot.v11.bot import Bot + from nonebot.adapters.onebot.v11.message import Message + + from nonebot_bison.plugin_config import plugin_config + from nonebot_bison.sub_manager import add_cookie_matcher + + mocker.patch.object(plugin_config, "bison_to_me", True) + + async with app.test_matcher(add_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event = fake_private_message_event(message=Message("添加cookie"), sender=fake_superuser) + ctx.receive_event(bot, event) + ctx.should_pass_rule() + ctx.should_pass_permission() + + async with app.test_matcher(add_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event = fake_private_message_event(message=Message("添加cookie"), sender=fake_admin_user) + ctx.receive_event(bot, event) + ctx.should_not_pass_rule() + ctx.should_pass_permission() + + +async def test_add_cookie_target_no_cookie(app: App, mocker: MockerFixture): + from nonebot.adapters.onebot.v11.bot import Bot + from nonebot.adapters.onebot.v11.message import Message + + from nonebot_bison.sub_manager import add_cookie_target_matcher + + async with app.test_matcher(add_cookie_target_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + from nonebug_saa import should_send_saa + from nonebot_plugin_saa import TargetQQGroup, MessageFactory + + from nonebot_bison.config import config + from nonebot_bison.types import Target as T_Target + + target = T_Target("weibo_id") + platform_name = "weibo" + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=target, + target_name="weibo_name", + platform_name=platform_name, + cats=[], + tags=[], + ) + + event_1 = fake_private_message_event( + message=Message("关联cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + should_send_saa( + ctx, + MessageFactory( + "订阅的帐号为：\n1 weibo weibo_name weibo_id\n []\n请输入要关联 cookie 的订阅的序号\n输入'取消'中止" + ), + bot, + event=event_1, + ) + event_2 = fake_private_message_event( + message=Message("1"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2) + ctx.should_pass_rule() + ctx.should_call_send( + event_2, + "当前平台暂无可关联的 Cookie，请使用“添加cookie”命令添加或检查已关联的 Cookie", + True, + ) + + +async def test_add_cookie(app: App, mocker: MockerFixture): + from nonebot.adapters.onebot.v11.bot import Bot + from nonebot.adapters.onebot.v11.message import Message + + from nonebot_bison.platform import platform_manager + from nonebot_bison.sub_manager import common_platform, add_cookie_matcher, add_cookie_target_matcher + + async with app.test_matcher(add_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event_1 = fake_private_message_event( + message=Message("添加Cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + ctx.should_call_send( + event_1, + BotReply.add_reply_on_add_cookie(platform_manager, common_platform), + True, + ) + event_2 = fake_private_message_event( + message=Message("全部"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2) + ctx.should_pass_rule() + ctx.should_rejected() + ctx.should_call_send( + event_2, + BotReply.add_reply_on_add_cookie_input_allplatform(platform_manager), + True, + ) + event_3 = fake_private_message_event( + message=Message("weibo"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_3) + ctx.should_pass_rule() + ctx.should_call_send(event_3, BotReply.add_reply_on_input_cookie) + event_4_err = fake_private_message_event( + message=Message("test"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_4_err) + ctx.should_call_send(event_4_err, "无效的 Cookie，请检查后重新输入，详情见<待添加的文档>", True) + ctx.should_rejected() + event_4_ok = fake_private_message_event( + message=Message(json.dumps({"cookie": "test"})), + sender=fake_superuser, + to_me=True, + user_id=fake_superuser.user_id, + ) + ctx.receive_event(bot, event_4_ok) + ctx.should_pass_rule() + ctx.should_call_send( + event_4_ok, '已添加 Cookie: {"cookie": "test"} 到平台 weibo\n请使用“关联cookie”为 Cookie 关联订阅', True + ) + + async with app.test_matcher(add_cookie_target_matcher) as ctx: + from nonebug_saa import should_send_saa + from nonebot_plugin_saa import TargetQQGroup, MessageFactory + + from nonebot_bison.config import config + from nonebot_bison.types import Target as T_Target + + target = T_Target("weibo_id") + platform_name = "weibo" + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=target, + target_name="weibo_name", + platform_name=platform_name, + cats=[], + tags=[], + ) + bot = ctx.create_bot(base=Bot) + event_1 = fake_private_message_event( + message=Message("关联cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + should_send_saa( + ctx, + MessageFactory( + "订阅的帐号为：\n1 weibo weibo_name weibo_id\n []\n请输入要关联 cookie 的订阅的序号\n输入'取消'中止" + ), + bot, + event=event_1, + ) + event_2_err = fake_private_message_event( + message=Message("2"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2_err) + ctx.should_call_send(event_2_err, "序号错误", True) + ctx.should_rejected() + event_2_ok = fake_private_message_event( + message=Message("1"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2_ok) + ctx.should_pass_rule() + ctx.should_call_send(event_2_ok, '请选择一个 Cookie，已关联的 Cookie 不会显示\n1. weibo.com [{"cookie":]', True) + event_3_err = fake_private_message_event( + message=Message("2"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_3_err) + ctx.should_call_send(event_3_err, "序号错误", True) + ctx.should_rejected() + event_3_ok = fake_private_message_event( + message=Message("1"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_3_ok) + ctx.should_pass_rule() + ctx.should_call_send(event_3_ok, '已关联 Cookie: weibo.com [{"cookie":] 到订阅 weibo.com weibo_id', True) + + +async def test_add_cookie_target_no_target(app: App, mocker: MockerFixture): + + from nonebot.adapters.onebot.v11.bot import Bot + from nonebot.adapters.onebot.v11.message import Message + + from nonebot_bison.sub_manager import add_cookie_target_matcher + + async with app.test_matcher(add_cookie_target_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event_1 = fake_private_message_event( + message=Message("关联cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + ctx.should_call_send( + event_1, + "暂无已订阅账号\n请使用“添加订阅”命令添加订阅", + True, + ) diff --git a/tests/utils.py b/tests/utils.py index 04efad5..24f459e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -89,6 +89,7 @@ add_reply_on_id_input_search = ( class BotReply: + @staticmethod def add_reply_on_platform(platform_manager, common_platform): return ( @@ -159,3 +160,33 @@ class BotReply: add_reply_on_tags_need_more_info = "订阅标签直接输入标签内容\n屏蔽标签请在标签名称前添加~号\n详见https://nonebot-bison.netlify.app/usage/#%E5%B9%B3%E5%8F%B0%E8%AE%A2%E9%98%85%E6%A0%87%E7%AD%BE-tag" add_reply_abort = "已中止订阅" no_permission = "您没有权限进行此操作，请联系 Bot 管理员" + + @staticmethod + def add_reply_on_add_cookie(platform_manager, common_platform): + from nonebot_bison.utils.site import is_cookie_client_manager + + return ( + "请输入想要添加 Cookie 的平台，目前支持，请输入冒号左边的名称：\n" + + "".join( + [ + f"{platform_name}: {platform_manager[platform_name].name}\n" + for platform_name in common_platform + if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr) + ] + ) + + "要查看全部平台请输入：“全部”\n中止添加cookie过程请输入：“取消”" + ) + + @staticmethod + def add_reply_on_add_cookie_input_allplatform(platform_manager): + from nonebot_bison.utils.site import is_cookie_client_manager + + return "全部平台\n" + "\n".join( + [ + f"{platform_name}: {platform.name}" + for platform_name, platform in platform_manager.items() + if is_cookie_client_manager(platform.site.client_mgr) + ] + ) + + add_reply_on_input_cookie = "请输入 Cookie" From c784417ecc9ad3499c5c8dbe42f66abe857ed136 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 13 Sep 2024 20:01:54 +0800 Subject: [PATCH 47/51] =?UTF-8?q?:white=5Fcheck=5Fmark:=20=E5=88=9D?= =?UTF-8?q?=E6=AD=A5=E5=AE=9E=E7=8E=B0=E5=88=A0=E9=99=A4cookie=E7=9A=84?= =?UTF-8?q?=E5=8D=95=E5=85=83=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/sub_manager/del_cookie.py | 4 +- tests/sub_manager/test_delete_cookie.py | 133 ++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 tests/sub_manager/test_delete_cookie.py diff --git a/nonebot_bison/sub_manager/del_cookie.py b/nonebot_bison/sub_manager/del_cookie.py index e343d4f..0ccbfe4 100644 --- a/nonebot_bison/sub_manager/del_cookie.py +++ b/nonebot_bison/sub_manager/del_cookie.py @@ -16,7 +16,7 @@ def do_del_cookie(del_cookie: type[Matcher]): async def send_list(state: T_State): cookies = await config.get_cookie(is_anonymous=False) if not cookies: - await del_cookie.finish("暂无已添加 Cookie\n请使用“添加cookie”命令添加") + await del_cookie.finish("暂无已添加的 Cookie\n请使用“添加cookie”命令添加") res = "已添加的 Cookie 为：\n" state["cookie_table"] = {} for index, cookie in enumerate(cookies, 1): @@ -40,6 +40,8 @@ def do_del_cookie(del_cookie: type[Matcher]): if cookie.targets: await del_cookie.reject("只能删除未关联的 Cookie，请使用“取消关联cookie”命令取消关联") await config.delete_cookie_by_id(cookie.id) + except KeyError: + await del_cookie.reject("序号错误") except Exception: await del_cookie.reject("删除错误") else: diff --git a/tests/sub_manager/test_delete_cookie.py b/tests/sub_manager/test_delete_cookie.py new file mode 100644 index 0000000..684739a --- /dev/null +++ b/tests/sub_manager/test_delete_cookie.py @@ -0,0 +1,133 @@ +import json + +from nonebug.app import App + +from ..utils import fake_superuser, fake_private_message_event + + +async def test_del_cookie_err(app: App): + from nonebug_saa import should_send_saa + from nonebot.adapters.onebot.v11.bot import Bot + from nonebot.adapters.onebot.v11.message import Message + from nonebot_plugin_saa import TargetQQGroup, MessageFactory + + from nonebot_bison.config import config + from nonebot_bison.config.db_model import Cookie + from nonebot_bison.types import Target as T_Target + from nonebot_bison.sub_manager import del_cookie_matcher + + async with app.test_matcher(del_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event = fake_private_message_event( + message=Message("删除cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event) + ctx.should_pass_rule() + ctx.should_pass_permission() + ctx.should_call_send(event, "暂无已添加的 Cookie\n请使用“添加cookie”命令添加", True) + + async with app.test_matcher(del_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + target = T_Target("weibo_id") + platform_name = "weibo" + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=target, + target_name="weibo_name", + platform_name=platform_name, + cats=[], + tags=[], + ) + await config.add_cookie(Cookie(content=json.dumps({"cookie": "test"}), site_name="weibo.com")) + cookies = await config.get_cookie(is_anonymous=False) + await config.add_cookie_target(target, platform_name, cookies[0].id) + + event_1 = fake_private_message_event( + message=Message("删除cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + ctx.should_pass_permission() + should_send_saa( + ctx, + MessageFactory( + '已添加的 Cookie 为：\n1 weibo.com weibo.com [{"cookie":] ' + "1个关联\n请输入要删除的 Cookie 的序号\n输入'取消'中止" + ), + bot, + event=event_1, + ) + event_2_err = fake_private_message_event( + message=Message("2"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2_err) + ctx.should_call_send(event_2_err, "序号错误", True) + ctx.should_rejected() + + event_2 = fake_private_message_event( + message=Message("1"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2) + ctx.should_pass_rule() + ctx.should_call_send(event_2, "只能删除未关联的 Cookie，请使用“取消关联cookie”命令取消关联", True) + ctx.should_call_send(event_2, "删除错误", True) + ctx.should_rejected() + + +async def test_del_cookie(app: App): + from nonebug_saa import should_send_saa + from nonebot.adapters.onebot.v11.bot import Bot + from nonebot.adapters.onebot.v11.message import Message + from nonebot_plugin_saa import TargetQQGroup, MessageFactory + + from nonebot_bison.config import config + from nonebot_bison.config.db_model import Cookie + from nonebot_bison.types import Target as T_Target + from nonebot_bison.sub_manager import del_cookie_matcher + + async with app.test_matcher(del_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event = fake_private_message_event( + message=Message("删除cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event) + ctx.should_pass_rule() + ctx.should_pass_permission() + ctx.should_call_send(event, "暂无已添加的 Cookie\n请使用“添加cookie”命令添加", True) + + async with app.test_matcher(del_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + target = T_Target("weibo_id") + platform_name = "weibo" + await config.add_subscribe( + TargetQQGroup(group_id=123), + target=target, + target_name="weibo_name", + platform_name=platform_name, + cats=[], + tags=[], + ) + await config.add_cookie(Cookie(content=json.dumps({"cookie": "test"}), site_name="weibo.com")) + + event_1 = fake_private_message_event( + message=Message("删除cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + ctx.should_pass_permission() + should_send_saa( + ctx, + MessageFactory( + '已添加的 Cookie 为：\n1 weibo.com weibo.com [{"cookie":]' + " 0个关联\n请输入要删除的 Cookie 的序号\n输入'取消'中止" + ), + bot, + event=event_1, + ) + event_2 = fake_private_message_event( + message=Message("1"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2) + ctx.should_pass_rule() + ctx.should_pass_permission() + ctx.should_call_send(event_2, "删除成功", True) From 76f271584f4af2f73a4513c47db981da6758057b Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 19 Sep 2024 20:28:16 +0800 Subject: [PATCH 48/51] =?UTF-8?q?:sparkles:=20web=20api=20=E5=88=9D?= =?UTF-8?q?=E7=A8=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/admin_page/api.py | 32 ++++++++++++++++++++++ nonebot_bison/admin_page/types.py | 45 +++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py index afe834e..e7b69d5 100644 --- a/nonebot_bison/admin_page/api.py +++ b/nonebot_bison/admin_page/api.py @@ -17,9 +17,11 @@ from .token_manager import token_manager from ..config.db_config import SubscribeDupException from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config from .types import ( + Cookie, TokenResp, GlobalConf, StatusResp, + CookieTarget, SubscribeResp, PlatformConfig, AddSubscribeReq, @@ -197,3 +199,33 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We except NoSuchTargetException: raise HTTPException(status.HTTP_400_BAD_REQUEST, "no such subscribe") return StatusResp(ok=True, msg="") + + +@router.get("/cookie") +async def get_cookie() -> list[Cookie]: + pass + + +@router.post("/cookie") +async def add_cookie(site_name: str, content: str) -> StatusResp: + pass + + +@router.delete("/cookie") +async def del_cookie(site_name: str, content: str) -> StatusResp: + pass + + +@router.get("/cookie_target") +async def get_cookie_target() -> list[CookieTarget]: + pass + + +@router.post("/cookie_target") +async def add_cookie_target(site_name: str, target: str) -> StatusResp: + pass + + +@router.delete("/cookie_target") +async def del_cookie_target(site_name: str, target: str) -> StatusResp: + pass diff --git a/nonebot_bison/admin_page/types.py b/nonebot_bison/admin_page/types.py index 7a18b67..ccbbd8e 100644 --- a/nonebot_bison/admin_page/types.py +++ b/nonebot_bison/admin_page/types.py @@ -6,14 +6,22 @@ class PlatformConfig(BaseModel): categories: dict[int, str] enabledTag: bool platformName: str + site_name: str hasTarget: bool +class SiteConfig(BaseModel): + name: str + enable_cookie: bool + + AllPlatformConf = dict[str, PlatformConfig] +AllSiteConf = dict[str, SiteConfig] class GlobalConf(BaseModel): platformConf: AllPlatformConf + siteConf: AllSiteConf class TokenResp(BaseModel): @@ -50,3 +58,40 @@ class AddSubscribeReq(BaseModel): class StatusResp(BaseModel): ok: bool msg: str + + +from typing import Any +from datetime import datetime + +from pydantic import BaseModel + + +class Target(BaseModel): + platform_name: str + target_name: str + target: str + + class Config: + orm_mode = True + + +class Cookie(BaseModel): + site_name: str + friendly_name: str + last_usage: datetime + status: str + cd_milliseconds: int + is_universal: bool + is_anonymous: bool + tags: dict[str, Any] + + class Config: + orm_mode = True + + +class CookieTarget(BaseModel): + target: Target + cookie: Cookie + + class Config: + orm_mode = True From d890d32bba039f9fda285c9d341bc7c75d294b49 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Thu, 19 Sep 2024 21:16:56 +0800 Subject: [PATCH 49/51] =?UTF-8?q?:sparkles:=20web=20api=20=E5=88=9D?= =?UTF-8?q?=E7=A8=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/admin_page/api.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py index e7b69d5..0ff0df4 100644 --- a/nonebot_bison/admin_page/api.py +++ b/nonebot_bison/admin_page/api.py @@ -201,31 +201,31 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We return StatusResp(ok=True, msg="") -@router.get("/cookie") +@router.get("/cookie", dependencies=[Depends(check_is_superuser)]) async def get_cookie() -> list[Cookie]: pass -@router.post("/cookie") +@router.post("/cookie", dependencies=[Depends(check_is_superuser)]) async def add_cookie(site_name: str, content: str) -> StatusResp: pass -@router.delete("/cookie") +@router.delete("/cookie", dependencies=[Depends(check_is_superuser)]) async def del_cookie(site_name: str, content: str) -> StatusResp: pass -@router.get("/cookie_target") +@router.get("/cookie_target", dependencies=[Depends(check_is_superuser)]) async def get_cookie_target() -> list[CookieTarget]: pass -@router.post("/cookie_target") +@router.post("/cookie_target", dependencies=[Depends(check_is_superuser)]) async def add_cookie_target(site_name: str, target: str) -> StatusResp: pass -@router.delete("/cookie_target") +@router.delete("/cookie_target", dependencies=[Depends(check_is_superuser)]) async def del_cookie_target(site_name: str, target: str) -> StatusResp: pass From c85e77c8011a9104d4307e78ddcd23c20f1c2c77 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 20 Sep 2024 00:25:40 +0800 Subject: [PATCH 50/51] =?UTF-8?q?:sparkles:=20web=20api=20=E6=9B=B4?= =?UTF-8?q?=E6=96=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/admin_page/api.py | 12 ++++++------ nonebot_bison/admin_page/types.py | 12 ++---------- 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py index 0ff0df4..b907813 100644 --- a/nonebot_bison/admin_page/api.py +++ b/nonebot_bison/admin_page/api.py @@ -202,7 +202,7 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We @router.get("/cookie", dependencies=[Depends(check_is_superuser)]) -async def get_cookie() -> list[Cookie]: +async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]: pass @@ -211,21 +211,21 @@ async def add_cookie(site_name: str, content: str) -> StatusResp: pass -@router.delete("/cookie", dependencies=[Depends(check_is_superuser)]) -async def del_cookie(site_name: str, content: str) -> StatusResp: +@router.delete("/cookie/{cookie_id}", dependencies=[Depends(check_is_superuser)]) +async def del_cookie(cookie_id: int) -> StatusResp: pass @router.get("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def get_cookie_target() -> list[CookieTarget]: +async def get_cookie_target(site_name: str = None, target: str = None, cookie_id: int = None) -> list[CookieTarget]: pass @router.post("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def add_cookie_target(site_name: str, target: str) -> StatusResp: +async def add_cookie_target(site_name: str, target: str, cookie_id: int) -> StatusResp: pass @router.delete("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def del_cookie_target(site_name: str, target: str) -> StatusResp: +async def del_cookie_target(site_name: str, target: str, cookie_id: int) -> StatusResp: pass diff --git a/nonebot_bison/admin_page/types.py b/nonebot_bison/admin_page/types.py index ccbbd8e..3c7ffd8 100644 --- a/nonebot_bison/admin_page/types.py +++ b/nonebot_bison/admin_page/types.py @@ -71,11 +71,9 @@ class Target(BaseModel): target_name: str target: str - class Config: - orm_mode = True - class Cookie(BaseModel): + id: int site_name: str friendly_name: str last_usage: datetime @@ -85,13 +83,7 @@ class Cookie(BaseModel): is_anonymous: bool tags: dict[str, Any] - class Config: - orm_mode = True - class CookieTarget(BaseModel): target: Target - cookie: Cookie - - class Config: - orm_mode = True + cookie_id: int From 94daf7435998440d84c3db213225f7dcbb82361b Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Fri, 20 Sep 2024 15:39:22 +0800 Subject: [PATCH 51/51] =?UTF-8?q?:sparkles:=20web=20api=20=E5=88=9D?= =?UTF-8?q?=E6=AD=A5=E5=AE=8C=E5=B7=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/admin_page/api.py | 67 ++++++++++++++++++++++++++------- 1 file changed, 53 insertions(+), 14 deletions(-) diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py index b907813..69b4ce2 100644 --- a/nonebot_bison/admin_page/api.py +++ b/nonebot_bison/admin_page/api.py @@ -1,3 +1,5 @@ +from typing import cast + import nonebot from fastapi import status from fastapi.routing import APIRouter @@ -12,14 +14,16 @@ from ..apis import check_sub_target from .jwt import load_jwt, pack_jwt from ..types import Target as T_Target from ..utils.get_bot import get_groups -from ..platform import platform_manager from .token_manager import token_manager from ..config.db_config import SubscribeDupException +from ..platform import site_manager, platform_manager +from ..utils.site import CookieClientManager, is_cookie_client_manager from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config from .types import ( Cookie, TokenResp, GlobalConf, + SiteConfig, StatusResp, CookieTarget, SubscribeResp, @@ -56,16 +60,20 @@ async def check_is_superuser(token_obj: dict = Depends(get_jwt_obj)): @router.get("/global_conf") async def get_global_conf() -> GlobalConf: - res = {} + platform_res = {} for platform_name, platform in platform_manager.items(): - res[platform_name] = PlatformConfig( + platform_res[platform_name] = PlatformConfig( platformName=platform_name, categories=platform.categories, enabledTag=platform.enable_tag, + site_name=platform.site.name, name=platform.name, hasTarget=getattr(platform, "has_target"), ) - return GlobalConf(platformConf=res) + site_res = {} + for site_name, site in site_manager.items(): + site_res[site_name] = SiteConfig(name=site_name, enable_cookie=is_cookie_client_manager(site.client_mgr)) + return GlobalConf(platformConf=platform_res, siteConf=site_res) async def get_admin_groups(qq: int): @@ -203,29 +211,60 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We @router.get("/cookie", dependencies=[Depends(check_is_superuser)]) async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]: - pass + cookies_in_db = await config.get_cookie(site_name, is_anonymous=False) + client_mgr = cast(CookieClientManager, site_manager[site_name].client_mgr) + friendly_names = [await client_mgr.get_cookie_friendly_name(x) for x in cookies_in_db] + return [ + Cookie( + id=cookies_in_db[i].id, + friendly_name=friendly_names[i], + site_name=cookies_in_db[i].site_name, + last_usage=cookies_in_db[i].last_usage, + status=cookies_in_db[i].status, + cd_milliseconds=cookies_in_db[i].cd_milliseconds, + is_universal=cookies_in_db[i].is_universal, + is_anonymous=cookies_in_db[i].is_anonymous, + tags=cookies_in_db[i].tags, + ) + for i in range(len(cookies_in_db)) + ] @router.post("/cookie", dependencies=[Depends(check_is_superuser)]) async def add_cookie(site_name: str, content: str) -> StatusResp: - pass + client_mgr = cast(CookieClientManager, site_manager[site_name].client_mgr) + await client_mgr.add_user_cookie(content) + return StatusResp(ok=True, msg="") @router.delete("/cookie/{cookie_id}", dependencies=[Depends(check_is_superuser)]) -async def del_cookie(cookie_id: int) -> StatusResp: - pass +async def delete_cookie_by_id(cookie_id: int) -> StatusResp: + await config.delete_cookie_by_id(cookie_id) + return StatusResp(ok=True, msg="") @router.get("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def get_cookie_target(site_name: str = None, target: str = None, cookie_id: int = None) -> list[CookieTarget]: - pass +async def get_cookie_target( + site_name: str | None = None, target: str | None = None, cookie_id: int | None = None +) -> list[CookieTarget]: + cookie_targets = await config.get_cookie_target() + # TODO: filter in SQL + return [ + x + for x in cookie_targets + if (site_name is None or x.cookie.site_name == site_name) + and (target is None or x.target.target == target) + and (cookie_id is None or x.cookie.id == cookie_id) + ] @router.post("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def add_cookie_target(site_name: str, target: str, cookie_id: int) -> StatusResp: - pass +async def add_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp: + await config.add_cookie_target(target, platform_name, cookie_id) + return StatusResp(ok=True, msg="") @router.delete("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def del_cookie_target(site_name: str, target: str, cookie_id: int) -> StatusResp: - pass +async def del_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp: + await config.delete_cookie_target(target, platform_name, cookie_id) + return StatusResp(ok=True, msg="")