✨ 添加 Cookie 组件 (#633)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

nonebot_bison/sub_manager/__init__.py

@@ -14,6 +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 .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info
 
 add_sub_matcher = on_command(
@@ -42,6 +46,46 @@ 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",
+    aliases={"添加Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_add_cookie(add_cookie_matcher)
+
+add_cookie_target_matcher = on_command(
+    "关联cookie",
+    aliases={"关联Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_add_cookie_target(add_cookie_target_matcher)
+
+del_cookie_target_matcher = on_command(
+    "取消关联cookie",
+    aliases={"取消关联Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_del_cookie_target(del_cookie_target_matcher)
+
+del_cookie_matcher = on_command(
+    "删除cookie",
+    aliases={"删除Cookie"},
+    rule=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, "已取消")
@@ -109,7 +153,11 @@ async def do_dispatch_command(
 
 
 no_permission_matcher = on_command(
-    "添加订阅", rule=configurable_to_me, aliases={"删除订阅", "群管理"}, priority=8, block=True
+    "添加订阅",
+    rule=configurable_to_me,
+    aliases={"删除订阅", "群管理", "管理后台", "添加cookie", "删除cookie", "关联cookie", "取消关联cookie"},
+    priority=8,
+    block=True,
 )
 
 
@@ -125,4 +173,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",
 ]

nonebot_bison/sub_manager/add_cookie.py

@@ -0,0 +1,82 @@
+from typing import cast
+from json import JSONDecodeError
+
+from nonebot.log import logger
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import Arg, ArgPlainText
+from nonebot.adapters.onebot.v11 import MessageEvent
+from nonebot.adapters import Message, MessageTemplate
+
+from ..scheduler import scheduler_dict
+from ..platform import platform_manager
+from ..utils.site import CookieClientManager, is_cookie_client_manager
+from .utils import common_platform, gen_handle_cancel, only_allow_private
+
+
+def do_add_cookie(add_cookie: type[Matcher]):
+    handle_cancel = gen_handle_cancel(add_cookie, "已中止添加cookie")
+
+    @add_cookie.handle()
+    async def init_promote(state: T_State, event: MessageEvent):
+        await only_allow_private(event, add_cookie)
+        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)
+                ]
+            )
+            + "要查看全部平台请输入：“全部”\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()
+                    if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr)
+                ]
+            )
+            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(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 = cast(CookieClientManager, scheduler_dict[platform_manager[state["platform"]].site].client_mgr)
+        cookie_text = cookie.extract_plain_text()
+        if not await client_mgr.validate_cookie(cookie_text):
+            await add_cookie.reject(
+                "无效的 Cookie，请检查后重新输入，详情见https://nonebot-bison.netlify.app/usage/cookie.html"
+            )
+        try:
+            cookie_name = await client_mgr.get_cookie_name(cookie_text)
+            state["cookie"] = cookie_text
+            state["cookie_name"] = cookie_name
+        except JSONDecodeError as e:
+            logger.error("获取 Cookie 名称失败" + str(e))
+            await add_cookie.reject(
+                "获取 Cookie 名称失败，请检查后重新输入，详情见https://nonebot-bison.netlify.app/usage/cookie.html"
+            )
+
+    @add_cookie.handle()
+    async def add_cookie_process(state: T_State):
+        client_mgr = cast(CookieClientManager, scheduler_dict[platform_manager[state["platform"]].site].client_mgr)
+        new_cookie = await client_mgr.add_user_cookie(state["cookie"], state["cookie_name"])
+        await add_cookie.finish(
+            f"已添加 Cookie: {new_cookie.cookie_name} 到平台 {state['platform']}"
+            + "\n请使用“关联cookie”为 Cookie 关联订阅"
+        )

nonebot_bison/sub_manager/add_cookie_target.py

@@ -0,0 +1,71 @@
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import ArgPlainText
+from nonebot_plugin_saa import MessageFactory
+from nonebot.adapters.onebot.v11 import MessageEvent
+from nonebot.internal.adapter import MessageTemplate
+
+from ..config import config
+from ..utils import parse_text
+from ..platform import platform_manager
+from .utils import gen_handle_cancel, only_allow_private, 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()
+    async def init_promote(state: T_State, event: MessageEvent):
+        await only_allow_private(event, add_cookie_target_matcher)
+        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()
+
+    @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]
+            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):
+
+        # 获取 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]
+        if not cookies:
+            await add_cookie_target_matcher.finish(
+                "当前平台暂无可关联的 Cookie，请使用“添加cookie”命令添加或检查已关联的 Cookie"
+            )
+        state["cookies"] = cookies
+
+        state["_prompt"] = "请选择一个 Cookie，已关联的 Cookie 不会显示\n" + "\n".join(
+            [f"{idx}. {cookie.cookie_name}" 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):
+        await config.add_cookie_target(state["target"]["target"], state["target"]["platform_name"], state["cookie"].id)
+        cookie = state["cookie"]
+        await add_cookie_target_matcher.finish(
+            f"已关联 Cookie: {cookie.cookie_name} " f"到订阅 {state['site'].name} {state['target']['target']}"
+        )

nonebot_bison/sub_manager/del_cookie.py

@@ -0,0 +1,47 @@
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import EventPlainText
+from nonebot_plugin_saa import MessageFactory
+from nonebot.adapters.onebot.v11 import MessageEvent
+
+from ..config import config
+from ..utils import parse_text
+from .utils import gen_handle_cancel, only_allow_private
+
+
+def do_del_cookie(del_cookie: type[Matcher]):
+    handle_cancel = gen_handle_cancel(del_cookie, "删除中止")
+
+    @del_cookie.handle()
+    async def send_list(state: T_State, event: MessageEvent):
+        await only_allow_private(event, del_cookie)
+        cookies = await config.get_cookie(is_anonymous=False)
+        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
+            res += f"{index} {cookie.site_name} {cookie.cookie_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 KeyError:
+            await del_cookie.reject("序号错误")
+        except Exception:
+            await del_cookie.reject("删除错误")
+        else:
+            await del_cookie.finish("删除成功")

nonebot_bison/sub_manager/del_cookie_target.py

@@ -0,0 +1,48 @@
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import EventPlainText
+from nonebot_plugin_saa import MessageFactory
+from nonebot.adapters.onebot.v11 import MessageEvent
+
+from ..config import config
+from ..utils import parse_text
+from .utils import gen_handle_cancel, only_allow_private
+
+
+def do_del_cookie_target(del_cookie_target: type[Matcher]):
+    handle_cancel = gen_handle_cancel(del_cookie_target, "取消关联中止")
+
+    @del_cookie_target.handle()
+    async def send_list(state: T_State, event: MessageEvent):
+        await only_allow_private(event, del_cookie_target)
+        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):
+            friendly_name = cookie_target.cookie.cookie_name
+            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("删除成功")

nonebot_bison/sub_manager/utils.py

@@ -1,16 +1,21 @@
 import contextlib
 from typing import Annotated
+from itertools import groupby
+from operator import attrgetter
 
 from nonebot.rule import Rule
 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 ..utils.site import is_cookie_client_manager
 
 
 def _configurable_to_me(to_me: bool = EventToMe()):
@@ -60,3 +65,67 @@ 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 | None = None,
+    is_index=False,
+    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:
+        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 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"
+    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(
+                    target=sub.target.target, site_name=platform.site.name, is_anonymous=False
+                )
+                if target_cookies:
+                    res += "  关联的 Cookie：\n"
+                    for cookie in target_cookies:
+                        res += f"  \t{cookie.cookie_name}\n"
+
+        else:
+            res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）"
+
+    return res
+
+
+async def only_allow_private(
+    event: Event,
+    matcher: type[Matcher],
+):
+    # if not issubclass(PrivateMessageEvent, event.__class__):
+    if event.message_type != "private":
+        await matcher.finish("请在私聊中使用此命令")