🐛 stash

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, "已取消")

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