✨ 导入导出支持 cookie

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,27 @@ 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(
+                site_name=cookie.site_name,
+                content=cookie.content,
+                cookie_name=cookie.cookie_name,
+                cd_milliseconds=cookie.cd_milliseconds,
+                is_universal=cookie.is_universal,
+                tags=cookie.tags,
+            )
+            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

@@ -11,11 +11,11 @@ from nonebot_plugin_datastore.db import create_session
 from sqlalchemy.orm.strategy_options import selectinload
 
 from .utils import NBESFVerMatchErr
-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,52 @@ 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 = 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)
+        cookies.append(
+            v3.Cookie(
+                site_name=cookie.site_name,
+                content=cookie.content,
+                cookie_name=cookie.cookie_name,
+                cd_milliseconds=cookie.cd_milliseconds,
+                is_universal=cookie.is_universal,
+                tags=cookie.tags,
+                targets=targets,
+            )
+        )
+
+    sub_group = v3.SubGroup(groups=groups, cookies=cookies)
 
     return sub_group
 
@@ -72,6 +102,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/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版本")