✨ 添加 Cookie 组件 (#633)

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

nonebot_bison/config/db_config.py

@@ -12,8 +12,8 @@ from nonebot_plugin_datastore import create_session
 
 from ..types import Tag
 from ..types import Target as T_Target
-from .utils import NoSuchTargetException
-from .db_model import User, Target, Subscribe, ScheduleTimeWeight
+from .utils import NoSuchTargetException, DuplicateCookieTargetException
+from .db_model import User, Cookie, Target, Subscribe, CookieTarget, ScheduleTimeWeight
 from ..types import Category, UserSubInfo, WeightConfig, TimeWeightConfig, PlatformWeightConfigResp
 
 
@@ -259,5 +259,125 @@ class DBConfig:
             )
         return res
 
+    async def get_cookie(
+        self,
+        site_name: str | None = None,
+        target: T_Target | None = None,
+        is_universal: bool | None = None,
+        is_anonymous: bool | None = None,
+    ) -> Sequence[Cookie]:
+        """获取满足传入条件的所有 cookie"""
+        async with create_session() as sess:
+            query = select(Cookie).distinct()
+            if is_universal is not None:
+                query = query.where(Cookie.is_universal == is_universal)
+            if is_anonymous is not None:
+                query = query.where(Cookie.is_anonymous == is_anonymous)
+            if site_name:
+                query = query.where(Cookie.site_name == site_name)
+            query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets))
+            res = (await sess.scalars(query)).all()
+            if target:
+                # 如果指定了 target，过滤掉不满足要求的cookie
+                query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target)
+                ids = set((await sess.scalars(query)).all())
+                # 如果指定了 target 且未指定 is_universal，则添加返回 universal cookie
+                res = [cookie for cookie in res if cookie.id in ids or cookie.is_universal]
+            return res
+
+    async def get_cookie_by_id(self, cookie_id: int) -> Cookie:
+        async with create_session() as sess:
+            cookie = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
+            return cookie
+
+    async def add_cookie(self, cookie: Cookie) -> int:
+        async with create_session() as sess:
+            sess.add(cookie)
+            await sess.commit()
+            await sess.refresh(cookie)
+            return cookie.id
+
+    async def update_cookie(self, cookie: Cookie):
+        async with create_session() as sess:
+            cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id))
+            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
+            await sess.commit()
+
+    async def delete_cookie_by_id(self, cookie_id: int):
+        async with create_session() as sess:
+            cookie = await sess.scalar(
+                select(Cookie)
+                .where(Cookie.id == cookie_id)
+                .outerjoin(CookieTarget)
+                .options(selectinload(Cookie.targets))
+            )
+            if len(cookie.targets) > 0:
+                raise Exception(f"cookie {cookie.id} in use")
+            await sess.execute(delete(Cookie).where(Cookie.id == cookie_id))
+            await sess.commit()
+
+    async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
+        """通过 cookie_id 可以唯一确定一个 Cookie，通过 target 和 platform_name 可以唯一确定一个 Target"""
+        async with create_session() as sess:
+            target_obj = await sess.scalar(
+                select(Target).where(Target.platform_name == platform_name, Target.target == target)
+            )
+            # check if relation exists
+            cookie_target = await sess.scalar(
+                select(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie_id == cookie_id)
+            )
+            if cookie_target:
+                raise DuplicateCookieTargetException()
+            cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
+            cookie_target = CookieTarget(target=target_obj, cookie=cookie_obj)
+            sess.add(cookie_target)
+            await sess.commit()
+
+    async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
+        async with create_session() as sess:
+            target_obj = await sess.scalar(
+                select(Target).where(Target.platform_name == platform_name, Target.target == target)
+            )
+            cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
+            await sess.execute(
+                delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj)
+            )
+            await sess.commit()
+
+    async def delete_cookie_target_by_id(self, cookie_target_id: int):
+        async with create_session() as sess:
+            await sess.execute(delete(CookieTarget).where(CookieTarget.id == cookie_target_id))
+            await sess.commit()
+
+    async def get_cookie_target(self) -> list[CookieTarget]:
+        async with create_session() as sess:
+            query = (
+                select(CookieTarget)
+                .outerjoin(Target)
+                .options(selectinload(CookieTarget.target))
+                .outerjoin(Cookie)
+                .options(selectinload(CookieTarget.cookie))
+            )
+            res = list((await sess.scalars(query)).all())
+            res.sort(key=lambda x: (x.target.platform_name, x.cookie_id, x.target_id))
+            return res
+
+    async def clear_db(self):
+        """清空数据库，用于单元测试清理环境"""
+        async with create_session() as sess:
+            await sess.execute(delete(User))
+            await sess.execute(delete(Target))
+            await sess.execute(delete(ScheduleTimeWeight))
+            await sess.execute(delete(Subscribe))
+            await sess.execute(delete(Cookie))
+            await sess.execute(delete(CookieTarget))
+            await sess.commit()
+
 
 config = DBConfig()

nonebot_bison/config/db_model.py

@@ -1,4 +1,5 @@
 import datetime
+from typing import Any
 from pathlib import Path
 
 from nonebot_plugin_saa import PlatformTarget
@@ -6,7 +7,7 @@ from sqlalchemy.dialects.postgresql import JSONB
 from nonebot.compat import PYDANTIC_V2, ConfigDict
 from nonebot_plugin_datastore import get_plugin_data
 from sqlalchemy.orm import Mapped, relationship, mapped_column
-from sqlalchemy import JSON, String, ForeignKey, UniqueConstraint
+from sqlalchemy import JSON, String, DateTime, ForeignKey, UniqueConstraint
 
 from ..types import Tag, Category
 
@@ -36,6 +37,7 @@ class Target(Model):
 
     subscribes: Mapped[list["Subscribe"]] = relationship(back_populates="target")
     time_weight: Mapped[list["ScheduleTimeWeight"]] = relationship(back_populates="target")
+    cookies: Mapped[list["CookieTarget"]] = relationship(back_populates="target")
 
 
 class ScheduleTimeWeight(Model):
@@ -66,3 +68,42 @@ class Subscribe(Model):
 
     target: Mapped[Target] = relationship(back_populates="subscribes")
     user: Mapped[User] = relationship(back_populates="subscribes")
+
+
+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 当前的状态
+    status: Mapped[str] = mapped_column(String(20), default="")
+    # 使用一次之后，需要的冷却时间
+    cd_milliseconds: Mapped[int] = mapped_column(default=0)
+    # 是否是通用 Cookie（对所有Target都有效）
+    is_universal: Mapped[bool] = mapped_column(default=False)
+    # 是否是匿名 Cookie
+    is_anonymous: Mapped[bool] = mapped_column(default=False)
+    # 标签，扩展用
+    tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={})
+
+    targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie")
+
+    @property
+    def cd(self) -> datetime.timedelta:
+        return datetime.timedelta(milliseconds=self.cd_milliseconds)
+
+    @cd.setter
+    def cd(self, value: datetime.timedelta):
+        self.cd_milliseconds = int(value.total_seconds() * 1000)
+
+
+class CookieTarget(Model):
+    id: Mapped[int] = mapped_column(primary_key=True)
+    target_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_target.id", ondelete="CASCADE"))
+    cookie_id: Mapped[int] = mapped_column(ForeignKey("nonebot_bison_cookie.id", ondelete="CASCADE"))
+
+    target: Mapped[Target] = relationship(back_populates="cookies")
+    cookie: Mapped[Cookie] = relationship(back_populates="targets")

nonebot_bison/config/migrations/f90b712557a9_add_cookie.py

@@ -0,0 +1,63 @@
+"""empty message
+
+Revision ID: f90b712557a9
+Revises: f9baef347cc8
+Create Date: 2024-09-23 10:03:30.593263
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import Text
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "f90b712557a9"
+down_revision = "f9baef347cc8"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "nonebot_bison_cookie",
+        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),
+        sa.Column("is_universal", sa.Boolean(), nullable=False),
+        sa.Column("is_anonymous", sa.Boolean(), nullable=False),
+        sa.Column("tags", sa.JSON().with_variant(postgresql.JSONB(astext_type=Text()), "postgresql"), nullable=False),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookie")),
+    )
+    op.create_table(
+        "nonebot_bison_cookietarget",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("target_id", sa.Integer(), nullable=False),
+        sa.Column("cookie_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["cookie_id"],
+            ["nonebot_bison_cookie.id"],
+            name=op.f("fk_nonebot_bison_cookietarget_cookie_id_nonebot_bison_cookie"),
+            ondelete="CASCADE",
+        ),
+        sa.ForeignKeyConstraint(
+            ["target_id"],
+            ["nonebot_bison_target.id"],
+            name=op.f("fk_nonebot_bison_cookietarget_target_id_nonebot_bison_target"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_nonebot_bison_cookietarget")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("nonebot_bison_cookietarget")
+    op.drop_table("nonebot_bison_cookie")
+    # ### end Alembic commands ###

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

@@ -0,0 +1,135 @@
+"""nbesf is Nonebot Bison Enchangable Subscribes File! ver.2"""
+
+from typing import Any
+from functools import partial
+
+from nonebot.log import logger
+from pydantic import BaseModel
+from nonebot_plugin_saa.registries import AllSupportedPlatformTarget
+from nonebot.compat import PYDANTIC_V2, ConfigDict, model_dump, type_validate_json, type_validate_python
+
+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 定义格式 ====== #
+NBESF_VERSION = 3
+
+
+class Target(BaseModel):
+    """Bsion快递包发货信息"""
+
+    target_name: str
+    target: str
+    platform_name: str
+    default_schedule_weight: int
+
+    if PYDANTIC_V2:
+        model_config = ConfigDict(from_attributes=True)
+    else:
+
+        class Config:
+            orm_mode = True
+
+
+class SubPayload(BaseModel):
+    """Bison快递包里的单件货物"""
+
+    categories: list[Category]
+    tags: list[Tag]
+    target: Target
+
+    if PYDANTIC_V2:
+        model_config = ConfigDict(from_attributes=True)
+    else:
+
+        class Config:
+            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给指定用户派送的快递包"""
+
+    # user_target: Bison快递包收货信息
+    user_target: AllSupportedPlatformTarget
+    subs: list[SubPayload]
+
+
+class SubGroup(NBESFBase):
+    """
+    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)
+
+        for sub in item.subs:
+            receipt = sub_receipt(
+                target=sub.target.target,
+                target_name=sub.target.target_name,
+                platform_name=sub.target.platform_name,
+                cats=sub.categories,
+                tags=sub.tags,
+            )
+            try:
+                await config.add_subscribe(receipt.user, **model_dump(receipt, exclude={"user"}))
+            except SubscribeDupException:
+                logger.warning(f"！添加订阅条目 {repr(receipt)} 失败: 相同的订阅已存在")
+            except Exception as e:
+                logger.error(f"！添加订阅条目 {repr(receipt)} 失败: {repr(e)}")
+            else:
+                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):
+            nbesf_data = type_validate_json(SubGroup, raw_data)
+        else:
+            nbesf_data = type_validate_python(SubGroup, raw_data)
+
+    except Exception as e:
+        logger.error("数据解析失败，该数据格式可能不满足NBESF格式标准！")
+        raise NBESFParseErr("数据解析失败") from e
+    else:
+        return nbesf_data

nonebot_bison/config/subs_io/subs_io.py

@@ -10,12 +10,13 @@ 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 ..db_model import User, Subscribe
-from .nbesf_model import NBESFBase, v1, v2
+from .. import config
+from .utils import NBESFVerMatchErr, row2dict
+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 +35,54 @@ 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: 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)
+
+    def cookie_transform(cookie: Cookie, targets: [Target]) -> v3.Cookie:
+        cookie_dict = row2dict(cookie)
+        cookie_dict["tags"] = cookie.tags
+        cookie_dict["targets"] = targets
+        return v3.Cookie(**cookie_dict)
+
+    cookies: list[v3.Cookie] = []
+    cookie_set = set()
+    for cookie, targets in cookie_target_dict.items():
+        assert isinstance(cookie, Cookie)
+        cookies.append(cookie_transform(cookie, targets))
+        cookie_set.add(cookie.id)
+
+    # 添加未关联的cookie
+    all_cookies = await config.get_cookie(is_anonymous=False)
+    cookies.extend([cookie_transform(c, []) for c in all_cookies if c.id not in cookie_set])
+
+    sub_group = v3.SubGroup(groups=groups, cookies=cookies)
 
     return sub_group
 
@@ -72,6 +105,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

@@ -1,4 +1,15 @@
+from ..db_model import Model
+
+
 class NBESFVerMatchErr(Exception): ...
 
 
 class NBESFParseErr(Exception): ...
+
+
+def row2dict(row: Model) -> dict:
+    d = {}
+    for column in row.__table__.columns:
+        d[column.name] = str(getattr(row, column.name))
+
+    return d

nonebot_bison/config/utils.py

@@ -8,3 +8,7 @@ class NoSuchSubscribeException(Exception):
 
 class NoSuchTargetException(Exception):
     pass
+
+
+class DuplicateCookieTargetException(Exception):
+    pass