From 6f20dbf3580dc76f92747f9311a28c0678b32b22 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 26 Aug 2024 17:28:59 +0800 Subject: [PATCH] =?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