Merge b04cbc2ced173fe90b20bf1eab7eabefd1006b1f into cf6b7fcd6dbc453af2cf5f2e98247afebd7b00d3

nonebot_bison/admin_page/api.py

@@ -210,15 +210,13 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We
 
 
 @router.get("/cookie", dependencies=[Depends(check_is_superuser)])
-async def get_cookie(site_name: str | None = None, target: str | None = None) -> list[Cookie]:
+async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]:
+    # todo: 调用 client_mgr 来添加cookie，以校验和获取cookie_name
     cookies_in_db = await config.get_cookie(site_name, is_anonymous=False)
-    # client_mgr = cast(CookieClientManager, site_manager[site_name].client_mgr)
-    # friendly_names = [await client_mgr.get_cookie_friendly_name(x) for x in cookies_in_db]
-    friendly_names = [x.content[:10] for x in cookies_in_db]
     return [
         Cookie(
             id=cookies_in_db[i].id,
-            friendly_name=friendly_names[i],
+            friendly_name=cookies_in_db[i].cookie_name,
             site_name=cookies_in_db[i].site_name,
             last_usage=cookies_in_db[i].last_usage,
             status=cookies_in_db[i].status,

nonebot_bison/config/db_config.py

@@ -298,6 +298,7 @@ class DBConfig:
             if not cookie_in_db:
                 raise ValueError(f"cookie {cookie.id} not found")
             cookie_in_db.content = cookie.content
+            cookie_in_db.cookie_name = cookie.cookie_name
             cookie_in_db.last_usage = cookie.last_usage
             cookie_in_db.status = cookie.status
             cookie_in_db.tags = cookie.tags

nonebot_bison/config/db_model.py

@@ -74,6 +74,8 @@ class Cookie(Model):
     id: Mapped[int] = mapped_column(primary_key=True)
     site_name: Mapped[str] = mapped_column(String(100))
     content: Mapped[str] = mapped_column(String(1024))
+    # Cookie 的友好名字，类似于 Target 的 target_name，用于展示
+    cookie_name: Mapped[str] = mapped_column(String(1024), default="unnamed cookie")
     # 最后使用的时刻
     last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1))
     # Cookie 当前的状态

nonebot_bison/config/migrations/f90b712557a9_add_cookie.py

@@ -1,8 +1,8 @@
 """empty message
 
-Revision ID: ef796b74b0fe
+Revision ID: f90b712557a9
 Revises: f9baef347cc8
-Create Date: 2024-09-13 00:34:08.601438
+Create Date: 2024-09-23 10:03:30.593263
 
 """
 
@@ -12,7 +12,7 @@ from sqlalchemy import Text
 from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
-revision = "ef796b74b0fe"
+revision = "f90b712557a9"
 down_revision = "f9baef347cc8"
 branch_labels = None
 depends_on = None
@@ -25,6 +25,7 @@ def upgrade() -> None:
         sa.Column("id", sa.Integer(), nullable=False),
         sa.Column("site_name", sa.String(length=100), nullable=False),
         sa.Column("content", sa.String(length=1024), nullable=False),
+        sa.Column("cookie_name", sa.String(length=1024), nullable=False),
         sa.Column("last_usage", sa.DateTime(), nullable=False),
         sa.Column("status", sa.String(length=20), nullable=False),
         sa.Column("cd_milliseconds", sa.Integer(), nullable=False),

nonebot_bison/config/subs_io/nbesf_model/__init__.py

@@ -1,6 +1,6 @@
 """nbesf is Nonebot Bison Enchangable Subscribes File!"""
 
-from . import v1, v2
+from . import v1, v2, v3
 from .base import NBESFBase
 
-__all__ = ["v1", "v2", "NBESFBase"]
+__all__ = ["v1", "v2", "v3", "NBESFBase"]

nonebot_bison/config/subs_io/nbesf_model/v3.py

@@ -11,6 +11,7 @@ from nonebot.compat import PYDANTIC_V2, ConfigDict, model_dump, type_validate_js
 from ..utils import NBESFParseErr
 from ....types import Tag, Category
 from .base import NBESFBase, SubReceipt
+from ...db_model import Cookie as DBCookie
 from ...db_config import SubscribeDupException, config
 
 # ===== nbesf 定义格式 ====== #
@@ -48,6 +49,18 @@ class SubPayload(BaseModel):
             orm_mode = True
 
 
+class Cookie(BaseModel):
+    """Bison的魔法饼干"""
+
+    site_name: str
+    content: str
+    cookie_name: str
+    cd_milliseconds: int
+    is_universal: bool
+    tags: dict[str, str]
+    targets: list[Target]
+
+
 class SubPack(BaseModel):
     """Bison给指定用户派送的快递包"""
 
@@ -58,19 +71,21 @@ class SubPack(BaseModel):
 
 class SubGroup(NBESFBase):
     """
-    Bison的全部订单(按用户分组)
+    Bison的全部订单(按用户分组)和魔法饼干
 
     结构参见`nbesf_model`下的对应版本
     """
 
     version: int = NBESF_VERSION
     groups: list[SubPack] = []
+    cookies: list[Cookie] = []
 
 
 # ======================= #
 
 
 async def subs_receipt_gen(nbesf_data: SubGroup):
+    logger.info("开始添加订阅流程")
     for item in nbesf_data.groups:
         sub_receipt = partial(SubReceipt, user=item.user_target)
 
@@ -92,6 +107,20 @@ async def subs_receipt_gen(nbesf_data: SubGroup):
                 logger.success(f"添加订阅条目 {repr(receipt)} 成功！")
 
 
+async def magic_cookie_gen(nbesf_data: SubGroup):
+    logger.info("开始添加 Cookie 流程")
+    for cookie in nbesf_data.cookies:
+        try:
+            new_cookie = DBCookie(**model_dump(cookie, exclude={"targets"}))
+            cookie_id = await config.add_cookie(new_cookie)
+            for target in cookie.targets:
+                await config.add_cookie_target(target.target, target.platform_name, cookie_id)
+        except Exception as e:
+            logger.error(f"！添加 Cookie 条目 {repr(cookie)} 失败: {repr(e)}")
+        else:
+            logger.success(f"添加 Cookie 条目 {repr(cookie)} 成功！")
+
+
 def nbesf_parser(raw_data: Any) -> SubGroup:
     try:
         if isinstance(raw_data, str):

nonebot_bison/config/subs_io/subs_io.py

@@ -10,12 +10,12 @@ from nonebot.compat import type_validate_python
 from nonebot_plugin_datastore.db import create_session
 from sqlalchemy.orm.strategy_options import selectinload
 
-from .utils import NBESFVerMatchErr
+from .utils import NBESFVerMatchErr, row2dict
-from ..db_model import User, Subscribe
+from .nbesf_model import NBESFBase, v1, v2, v3
-from .nbesf_model import NBESFBase, v1, v2
+from ..db_model import User, Cookie, Target, Subscribe, CookieTarget
 
 
-async def subscribes_export(selector: Callable[[Select], Select]) -> v2.SubGroup:
+async def subscribes_export(selector: Callable[[Select], Select]) -> v3.SubGroup:
     """
     将Bison订阅导出为 Nonebot Bison Exchangable Subscribes File 标准格式的 SubGroup 类型数据
 
@@ -34,22 +34,45 @@ async def subscribes_export(selector: Callable[[Select], Select]) -> v2.SubGroup
         user_stmt = cast(Select[tuple[User]], user_stmt)
         user_data = await sess.scalars(user_stmt)
 
-    groups: list[v2.SubPack] = []
+    groups: list[v3.SubPack] = []
-    user_id_sub_dict: dict[int, list[v2.SubPayload]] = defaultdict(list)
+    user_id_sub_dict: dict[int, list[v3.SubPayload]] = defaultdict(list)
 
     for sub in sub_data:
-        sub_paylaod = type_validate_python(v2.SubPayload, sub)
+        sub_paylaod = type_validate_python(v3.SubPayload, sub)
         user_id_sub_dict[sub.user_id].append(sub_paylaod)
 
     for user in user_data:
         assert isinstance(user, User)
-        sub_pack = v2.SubPack(
+        sub_pack = v3.SubPack(
             user_target=PlatformTarget.deserialize(user.user_target),
             subs=user_id_sub_dict[user.id],
         )
         groups.append(sub_pack)
 
-    sub_group = v2.SubGroup(groups=groups)
+    async with create_session() as sess:
+        cookie_target_stmt = (
+            select(CookieTarget)
+            .join(Cookie)
+            .join(Target)
+            .options(selectinload(CookieTarget.target))
+            .options(selectinload(CookieTarget.cookie))
+        )
+        cookie_target_data = await sess.scalars(cookie_target_stmt)
+
+    cookie_target_dict: dict[Cookie, list[v3.Target]] = defaultdict(list)
+    for cookie_target in cookie_target_data:
+        target_payload = type_validate_python(v3.Target, cookie_target.target)
+        cookie_target_dict[cookie_target.cookie].append(target_payload)
+
+    cookies: list[v3.Cookie] = []
+    for cookie, targets in cookie_target_dict.items():
+        assert isinstance(cookie, Cookie)
+        cookie_dict = row2dict(cookie)
+        cookie_dict["tags"] = cookie.tags
+        cookie_dict["targets"] = targets
+        cookies.append(v3.Cookie(**cookie_dict))
+
+    sub_group = v3.SubGroup(groups=groups, cookies=cookies)
 
     return sub_group
 
@@ -72,6 +95,10 @@ async def subscribes_import(
         case 2:
             assert isinstance(nbesf_data, v2.SubGroup)
             await v2.subs_receipt_gen(nbesf_data)
+        case 3:
+            assert isinstance(nbesf_data, v3.SubGroup)
+            await v3.subs_receipt_gen(nbesf_data)
+            await v3.magic_cookie_gen(nbesf_data)
         case _:
             raise NBESFVerMatchErr(f"不支持的NBESF版本：{nbesf_data.version}")
     logger.info("订阅流程结束，请检查所有订阅记录是否全部添加成功")

nonebot_bison/config/subs_io/utils.py

@@ -2,3 +2,11 @@ class NBESFVerMatchErr(Exception): ...
 
 
 class NBESFParseErr(Exception): ...
+
+
+def row2dict(row):
+    d = {}
+    for column in row.__table__.columns:
+        d[column.name] = str(getattr(row, column.name))
+
+    return d

nonebot_bison/platform/rss.py

@@ -9,11 +9,11 @@ from bs4 import BeautifulSoup as bs
 from ..post import Post
 from .platform import NewMessage
 from ..types import Target, RawPost
-from ..utils import Site, text_similarity
+from ..utils import text_similarity
-from ..utils.site import create_cookie_client_manager
+from ..utils.site import CookieSite, create_cookie_client_manager
 
 
-class RssSite(Site):
+class RssSite(CookieSite):
     name = "rss"
     schedule_type = "interval"
     schedule_setting = {"seconds": 30}

nonebot_bison/platform/weibo.py

@@ -11,10 +11,10 @@ from nonebot.log import logger
 from bs4 import BeautifulSoup as bs
 
 from ..post import Post
+from ..utils import http_client
 from .platform import NewMessage
-from ..utils import Site, http_client
-from ..utils.site import create_cookie_client_manager
 from ..types import Tag, Target, RawPost, ApiError, Category
+from ..utils.site import CookieSite, create_cookie_client_manager
 
 _HEADER = {
     "accept": (
@@ -36,7 +36,7 @@ _HEADER = {
 }
 
 
-class WeiboSite(Site):
+class WeiboSite(CookieSite):
     name = "weibo.com"
     schedule_type = "interval"
     schedule_setting = {"seconds": 3}

nonebot_bison/script/cli.py

@@ -11,7 +11,7 @@ from nonebot.log import logger
 from nonebot.compat import model_dump
 
 from ..scheduler.manager import init_scheduler
-from ..config.subs_io.nbesf_model import v1, v2
+from ..config.subs_io.nbesf_model import v1, v2, v3
 from ..config.subs_io import subscribes_export, subscribes_import
 
 try:
@@ -151,6 +151,8 @@ async def subs_import(path: str, format: str):
                 nbesf_data = v1.nbesf_parser(import_items)
             case 2:
                 nbesf_data = v2.nbesf_parser(import_items)
+            case 3:
+                nbesf_data = v3.nbesf_parser(import_items)
             case _:
                 raise NotImplementedError("不支持的NBESF版本")
 

nonebot_bison/sub_manager/add_cookie_target.py

@@ -1,5 +1,3 @@
-from typing import cast
-
 from nonebot.typing import T_State
 from nonebot.matcher import Matcher
 from nonebot.params import ArgPlainText
@@ -9,7 +7,6 @@ from nonebot.internal.adapter import MessageTemplate
 from ..config import config
 from ..utils import parse_text
 from ..platform import platform_manager
-from ..utils.site import CookieClientManager
 from .utils import gen_handle_cancel, generate_sub_list_text
 
 
@@ -51,9 +48,8 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]):
             )
         state["cookies"] = cookies
 
-        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)]
+            [f"{idx}. {cookie.cookie_name}" for idx, cookie in enumerate(cookies, 1)]
         )
 
     @add_cookie_target_matcher.got("cookie_idx", MessageTemplate("{_prompt}"), [handle_cancel])
@@ -68,8 +64,6 @@ def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]):
     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"]
-        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"已关联 Cookie: {cookie.cookie_name} " f"到订阅 {state['site'].name} {state['target']['target']}"
-            f"到订阅 {state['site'].name} {state['target']['target']}"
         )

nonebot_bison/sub_manager/del_cookie.py

@@ -5,7 +5,6 @@ from nonebot_plugin_saa import MessageFactory
 
 from ..config import config
 from ..utils import parse_text
-from ..platform import site_manager
 from .utils import gen_handle_cancel
 
 
@@ -21,9 +20,7 @@ def do_del_cookie(del_cookie: type[Matcher]):
         state["cookie_table"] = {}
         for index, cookie in enumerate(cookies, 1):
             state["cookie_table"][index] = cookie
-            client_mgr = site_manager[cookie.site_name].client_mgr
+            res += f"{index} {cookie.site_name} {cookie.cookie_name} {len(cookie.targets)}个关联\n"
-            friendly_name = await client_mgr.get_cookie_friendly_name(cookie)
-            res += f"{index} {cookie.site_name} {friendly_name} {len(cookie.targets)}个关联\n"
             if res[-1] != "\n":
                 res += "\n"
         res += "请输入要删除的 Cookie 的序号\n输入'取消'中止"

nonebot_bison/sub_manager/del_cookie_target.py

@@ -1,5 +1,3 @@
-from typing import cast
-
 from nonebot.typing import T_State
 from nonebot.matcher import Matcher
 from nonebot.params import EventPlainText
@@ -8,8 +6,6 @@ from nonebot_plugin_saa import MessageFactory
 from ..config import config
 from ..utils import parse_text
 from .utils import gen_handle_cancel
-from ..platform import platform_manager
-from ..utils.site import CookieClientManager
 
 
 def do_del_cookie_target(del_cookie_target: type[Matcher]):
@@ -23,8 +19,7 @@ def do_del_cookie_target(del_cookie_target: type[Matcher]):
         res = "已关联的 Cookie 为：\n"
         state["cookie_target_table"] = {}
         for index, cookie_target in enumerate(cookie_targets, 1):
-            client_mgr = cast(CookieClientManager, platform_manager[cookie_target.target.platform_name].site.client_mgr)
+            friendly_name = cookie_target.cookie.cookie_name
-            friendly_name = await client_mgr.get_cookie_friendly_name(cookie_target.cookie)
             state["cookie_target_table"][index] = {
                 "platform_name": cookie_target.target.platform_name,
                 "target": cookie_target.target,

nonebot_bison/sub_manager/utils.py

@@ -1,7 +1,7 @@
 import contextlib
+from typing import Annotated
 from itertools import groupby
 from operator import attrgetter
-from typing import Annotated, cast
 
 from nonebot.rule import Rule
 from nonebot.adapters import Event
@@ -13,9 +13,9 @@ 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 ..platform import site_manager, platform_manager
+from ..utils.site import is_cookie_client_manager
-from ..utils.site import CookieClientManager, is_cookie_client_manager
 
 
 def _configurable_to_me(to_me: bool = EventToMe()):
@@ -114,8 +114,7 @@ async def generate_sub_list_text(
                 if target_cookies:
                     res += "  关联的 Cookie：\n"
                     for cookie in target_cookies:
-                        client_mgr = cast(CookieClientManager, site_manager[platform.site.name].client_mgr)
+                        res += f"  \t{cookie.cookie_name}\n"
-                        res += f"  \t{await client_mgr.get_cookie_friendly_name(cookie)}\n"
 
         else:
             res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）"

nonebot_bison/utils/site.py

@@ -1,6 +1,6 @@
 import json
-from typing import Literal
 from json import JSONDecodeError
+from typing import Literal, cast
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
 
@@ -62,13 +62,20 @@ class CookieClientManager(ClientManager):
     @classmethod
     async def add_user_cookie(cls, content: str):
         """添加用户 cookie"""
+        from ..platform import site_manager
+
+        cookie_site = cast(type[CookieSite], site_manager[cls._site_name])
+        if not await cls.validate_cookie(content):
+            raise ValueError()
         cookie = Cookie(site_name=cls._site_name, content=content)
+        cookie.cookie_name = cookie_site.get_cookie_name(content)
         cookie.cd = cls._default_cd
         await config.add_cookie(cookie)
 
     @classmethod
     async def validate_cookie(cls, content: str) -> bool:
         """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写"""
+        # todo: 考虑移动到 cookie site 中
         try:
             data = json.loads(content)
             if not isinstance(data, dict):
@@ -77,13 +84,6 @@ class CookieClientManager(ClientManager):
             return False
         return True
 
-    @classmethod
-    async def get_cookie_friendly_name(cls, cookie: Cookie) -> str:
-        """获取 cookie 的友好名字，用于展示"""
-        from . import text_fletten
-
-        return text_fletten(f"{cookie.site_name} [{cookie.content[:10]}]")
-
     def _generate_hook(self, cookie: Cookie) -> callable:
         """hook 函数生成器，用于回写请求状态到数据库"""
 
@@ -156,12 +156,23 @@ class Site(metaclass=RegistryMeta, base=True):
     client_mgr: type[ClientManager] = DefaultClientManager
     require_browser: bool = False
     registry: list[type["Site"]]
-    cookie_format_prompt = "无效的 Cookie，请检查后重新输入，详情见<待添加的文档>"
 
     def __str__(self):
         return f"[{self.name}]-{self.name}-{self.schedule_setting}"
 
 
+class CookieSite(Site):
+    client_mgr: type[CookieClientManager] = CookieClientManager
+    cookie_format_prompt = "无效的 Cookie，请检查后重新输入，详情见<待添加的文档>"
+
+    @classmethod
+    def get_cookie_name(cls, content: str) -> str:
+        """从cookie内容中获取cookie的友好名字，添加cookie时调用，持久化在数据库中"""
+        from . import text_fletten
+
+        return text_fletten(f"{cls.name} [{content[:10]}]")
+
+
 def anonymous_site(schedule_type: Literal["date", "interval", "cron"], schedule_setting: dict) -> type[Site]:
     return type(
         "AnonymousSite",