:recycles: add_cookie_target

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(

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):

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(
+add_cookie_target_matcher = on_command(
-#     "关联cookie",
+    "关联cookie",
-#     rule=configurable_to_me,
+    rule=configurable_to_me,
-#     permission=SUPERUSER,
+    permission=SUPERUSER,
-#     priority=5,
+    priority=5,
-#     block=True,
+    block=True,
-# )
+)
-# do_add_cookie_target(add_cookie_target_matcher)
+do_add_cookie_target(add_cookie_target_matcher)
-#
+
 # del_cookie_target_matcher = on_command(
 #     "取消关联cookie",
 #     rule=configurable_to_me,

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']}"
         )

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: