From 404b1e445cfd494a690149bfa815aaba731b0b35 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 9 Sep 2024 18:23:23 +0800 Subject: [PATCH] :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: