From 5d160dcadcd4be828af3ba9c4502b2018d045247 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Mon, 23 Sep 2024 11:44:33 +0800 Subject: [PATCH] =?UTF-8?q?:sparkles:=20=E5=AF=BC=E5=85=A5=E5=AF=BC?= =?UTF-8?q?=E5=87=BA=E6=94=AF=E6=8C=81=20cookie?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../config/subs_io/nbesf_model/__init__.py | 4 +- .../config/subs_io/nbesf_model/v3.py | 38 +++++++++++++- nonebot_bison/config/subs_io/subs_io.py | 50 ++++++++++++++++--- nonebot_bison/script/cli.py | 4 +- 4 files changed, 84 insertions(+), 12 deletions(-) diff --git a/nonebot_bison/config/subs_io/nbesf_model/__init__.py b/nonebot_bison/config/subs_io/nbesf_model/__init__.py index 8dae14f..dff07d0 100644 --- a/nonebot_bison/config/subs_io/nbesf_model/__init__.py +++ b/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"] diff --git a/nonebot_bison/config/subs_io/nbesf_model/v3.py b/nonebot_bison/config/subs_io/nbesf_model/v3.py index 6f525e3..846c4ff 100644 --- a/nonebot_bison/config/subs_io/nbesf_model/v3.py +++ b/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): diff --git a/nonebot_bison/config/subs_io/subs_io.py b/nonebot_bison/config/subs_io/subs_io.py index ec82695..779b6c9 100644 --- a/nonebot_bison/config/subs_io/subs_io.py +++ b/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 +from .nbesf_model import NBESFBase, v1, v2, v3 +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] = [] - user_id_sub_dict: dict[int, list[v2.SubPayload]] = defaultdict(list) + groups: list[v3.SubPack] = [] + 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("订阅流程结束，请检查所有订阅记录是否全部添加成功") diff --git a/nonebot_bison/script/cli.py b/nonebot_bison/script/cli.py index 6240461..790a7dd 100644 --- a/nonebot_bison/script/cli.py +++ b/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版本")