✨ 添加 Cookie 组件 (#633)

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

nonebot_bison/admin_page/__init__.py

@@ -61,7 +61,7 @@ def init_fastapi(driver: "Driver"):
 
 
 def register_get_token_handler():
-    get_token = on_command("后台管理", rule=to_me(), priority=5, aliases={"管理后台"})
+    get_token = on_command("后台管理", rule=to_me(), priority=5, aliases={"管理后台"}, block=True)
 
     @get_token.handle()
     async def send_token(bot: "Bot", event: PrivateMessageEvent, state: T_State):

nonebot_bison/admin_page/api.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 import nonebot
 from fastapi import status
 from fastapi.routing import APIRouter
@@ -10,16 +12,21 @@ from fastapi.security.oauth2 import OAuth2PasswordBearer
 from ..types import WeightConfig
 from ..apis import check_sub_target
 from .jwt import load_jwt, pack_jwt
+from ..scheduler import scheduler_dict
 from ..types import Target as T_Target
 from ..utils.get_bot import get_groups
 from ..platform import platform_manager
 from .token_manager import token_manager
 from ..config.db_config import SubscribeDupException
+from ..utils.site import CookieClientManager, site_manager, is_cookie_client_manager
 from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config
 from .types import (
+    Cookie,
     TokenResp,
     GlobalConf,
+    SiteConfig,
     StatusResp,
+    CookieTarget,
     SubscribeResp,
     PlatformConfig,
     AddSubscribeReq,
@@ -54,16 +61,20 @@ async def check_is_superuser(token_obj: dict = Depends(get_jwt_obj)):
 
 @router.get("/global_conf")
 async def get_global_conf() -> GlobalConf:
-    res = {}
+    platform_res = {}
     for platform_name, platform in platform_manager.items():
-        res[platform_name] = PlatformConfig(
+        platform_res[platform_name] = PlatformConfig(
             platformName=platform_name,
             categories=platform.categories,
             enabledTag=platform.enable_tag,
+            siteName=platform.site.name,
             name=platform.name,
             hasTarget=getattr(platform, "has_target"),
         )
-    return GlobalConf(platformConf=res)
+    site_res = {}
+    for site_name, site in site_manager.items():
+        site_res[site_name] = SiteConfig(name=site_name, enable_cookie=is_cookie_client_manager(site.client_mgr))
+    return GlobalConf(platformConf=platform_res, siteConf=site_res)
 
 
 async def get_admin_groups(qq: int):
@@ -197,3 +208,72 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We
     except NoSuchTargetException:
         raise HTTPException(status.HTTP_400_BAD_REQUEST, "no such subscribe")
     return StatusResp(ok=True, msg="")
+
+
+@router.get("/cookie", dependencies=[Depends(check_is_superuser)])
+async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]:
+    cookies_in_db = await config.get_cookie(site_name, is_anonymous=False)
+    return [
+        Cookie(
+            id=cookies_in_db[i].id,
+            content=cookies_in_db[i].content,
+            cookie_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,
+            cd_milliseconds=cookies_in_db[i].cd_milliseconds,
+            is_universal=cookies_in_db[i].is_universal,
+            is_anonymous=cookies_in_db[i].is_anonymous,
+            tags=cookies_in_db[i].tags,
+        )
+        for i in range(len(cookies_in_db))
+    ]
+
+
+@router.post("/cookie", dependencies=[Depends(check_is_superuser)])
+async def add_cookie(site_name: str, content: str) -> StatusResp:
+    client_mgr = cast(CookieClientManager, scheduler_dict[site_manager[site_name]].client_mgr)
+    await client_mgr.add_user_cookie(content)
+    return StatusResp(ok=True, msg="")
+
+
+@router.delete("/cookie/{cookie_id}", dependencies=[Depends(check_is_superuser)])
+async def delete_cookie_by_id(cookie_id: int) -> StatusResp:
+    await config.delete_cookie_by_id(cookie_id)
+    return StatusResp(ok=True, msg="")
+
+
+@router.get("/cookie_target", dependencies=[Depends(check_is_superuser)])
+async def get_cookie_target(
+    site_name: str | None = None, target: str | None = None, cookie_id: int | None = None
+) -> list[CookieTarget]:
+    cookie_targets = await config.get_cookie_target()
+    # TODO: filter in SQL
+    return [
+        x
+        for x in cookie_targets
+        if (site_name is None or x.cookie.site_name == site_name)
+        and (target is None or x.target.target == target)
+        and (cookie_id is None or x.cookie.id == cookie_id)
+    ]
+
+
+@router.post("/cookie_target", dependencies=[Depends(check_is_superuser)])
+async def add_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp:
+    await config.add_cookie_target(target, platform_name, cookie_id)
+    return StatusResp(ok=True, msg="")
+
+
+@router.delete("/cookie_target", dependencies=[Depends(check_is_superuser)])
+async def del_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp:
+    await config.delete_cookie_target(target, platform_name, cookie_id)
+    return StatusResp(ok=True, msg="")
+
+
+@router.post("/cookie/validate", dependencies=[Depends(check_is_superuser)])
+async def get_cookie_valid(site_name: str, content: str) -> StatusResp:
+    client_mgr = cast(CookieClientManager, scheduler_dict[site_manager[site_name]].client_mgr)
+    if await client_mgr.validate_cookie(content):
+        return StatusResp(ok=True, msg="")
+    else:
+        return StatusResp(ok=False, msg="")

nonebot_bison/admin_page/types.py

@@ -6,14 +6,22 @@ class PlatformConfig(BaseModel):
     categories: dict[int, str]
     enabledTag: bool
     platformName: str
+    siteName: str
     hasTarget: bool
 
 
+class SiteConfig(BaseModel):
+    name: str
+    enable_cookie: bool
+
+
 AllPlatformConf = dict[str, PlatformConfig]
+AllSiteConf = dict[str, SiteConfig]
 
 
 class GlobalConf(BaseModel):
     platformConf: AllPlatformConf
+    siteConf: AllSiteConf
 
 
 class TokenResp(BaseModel):
@@ -50,3 +58,33 @@ class AddSubscribeReq(BaseModel):
 class StatusResp(BaseModel):
     ok: bool
     msg: str
+
+
+from typing import Any
+from datetime import datetime
+
+from pydantic import BaseModel
+
+
+class Target(BaseModel):
+    platform_name: str
+    target_name: str
+    target: str
+
+
+class Cookie(BaseModel):
+    id: int
+    site_name: str
+    content: str
+    cookie_name: str
+    last_usage: datetime
+    status: str
+    cd_milliseconds: int
+    is_universal: bool
+    is_anonymous: bool
+    tags: dict[str, Any]
+
+
+class CookieTarget(BaseModel):
+    target: Target
+    cookie_id: int

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

nonebot_bison/platform/bilibili/scheduler.py

@@ -1,13 +1,17 @@
+import json
 import random
 from typing_extensions import override
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, TypeVar
 
 from httpx import AsyncClient
 from nonebot import logger, require
 from playwright.async_api import Cookie
 
-from nonebot_bison.types import Target
-from nonebot_bison.utils import Site, ClientManager, http_client
+from nonebot_bison.utils import Site, http_client
+
+from ...utils.site import CookieClientManager
+from ...config.db_model import Cookie as CookieModel
 
 if TYPE_CHECKING:
     from .platforms import Bilibili
@@ -18,12 +22,9 @@ from nonebot_plugin_htmlrender import get_browser
 B = TypeVar("B", bound="Bilibili")
 
 
-class BilibiliClientManager(ClientManager):
-    _client: AsyncClient
-    _inited: bool = False
+class BilibiliClientManager(CookieClientManager):
 
-    def __init__(self) -> None:
-        self._client = http_client()
+    _default_cookie_cd = timedelta(seconds=120)
 
     async def _get_cookies(self) -> list[Cookie]:
         browser = await get_browser()
@@ -36,29 +37,33 @@ class BilibiliClientManager(ClientManager):
 
         return cookies
 
-    async def _reset_client_cookies(self, cookies: list[Cookie]):
+    def _gen_json_cookie(self, cookies: list[Cookie]):
+        cookie_dict = {}
         for cookie in cookies:
-            self._client.cookies.set(
-                name=cookie.get("name", ""),
-                value=cookie.get("value", ""),
-                domain=cookie.get("domain", ""),
-                path=cookie.get("path", "/"),
-            )
+            cookie_dict[cookie.get("name", "")] = cookie.get("value", "")
+        return cookie_dict
+
+    @override
+    async def _generate_anonymous_cookie(self) -> CookieModel:
+        cookies = await self._get_cookies()
+        cookie = CookieModel(
+            cookie_name=f"{self._site_name} anonymous",
+            site_name=self._site_name,
+            content=json.dumps(self._gen_json_cookie(cookies)),
+            is_universal=True,
+            is_anonymous=True,
+            last_usage=datetime.now(),
+            cd_milliseconds=0,
+            tags="{}",
+            status="",
+        )
+        return cookie
 
     @override
     async def refresh_client(self):
-        cookies = await self._get_cookies()
-        await self._reset_client_cookies(cookies)
+        await self._refresh_anonymous_cookie()
         logger.debug("刷新B站客户端的cookie")
 
-    @override
-    async def get_client(self, target: Target | None) -> AsyncClient:
-        if not self._inited:
-            logger.debug("初始化B站客户端")
-            await self.refresh_client()
-            self._inited = True
-        return self._client
-
     @override
     async def get_client_for_static(self) -> AsyncClient:
         return http_client()

nonebot_bison/platform/rss.py

@@ -9,13 +9,15 @@ 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 Site, CookieClientManager
 
 
 class RssSite(Site):
     name = "rss"
     schedule_type = "interval"
     schedule_setting = {"seconds": 30}
+    client_mgr = CookieClientManager.from_name(name)
 
 
 class RssPost(Post):
@@ -63,7 +65,7 @@ class Rss(NewMessage):
         return post.id
 
     async def get_sub_list(self, target: Target) -> list[RawPost]:
-        client = await self.ctx.get_client()
+        client = await self.ctx.get_client(target)
         res = await client.get(target, timeout=10.0)
         feed = feedparser.parse(res)
         entries = feed.entries

nonebot_bison/platform/weibo.py

@@ -3,6 +3,7 @@ import json
 from typing import Any
 from datetime import datetime
 from urllib.parse import unquote
+from typing_extensions import override
 
 from yarl import URL
 from lxml.etree import HTML
@@ -12,7 +13,8 @@ from bs4 import BeautifulSoup as bs
 
 from ..post import Post
 from .platform import NewMessage
-from ..utils import Site, http_client
+from ..utils import http_client, text_fletten
+from ..utils.site import Site, CookieClientManager
 from ..types import Tag, Target, RawPost, ApiError, Category
 
 _HEADER = {
@@ -35,10 +37,30 @@ _HEADER = {
 }
 
 
+class WeiboClientManager(CookieClientManager):
+    _site_name = "weibo.com"
+
+    async def _get_current_user_name(self, cookies: dict) -> str:
+        url = "https://m.weibo.cn/setup/nick/detail"
+        async with http_client() as client:
+            r = await client.get(url, headers=_HEADER, cookies=cookies)
+            data = json.loads(r.text)
+            name = data["data"]["user"]["screen_name"]
+            return name
+
+    @override
+    async def get_cookie_name(self, content: str) -> str:
+        """从cookie内容中获取cookie的友好名字，添加cookie时调用，持久化在数据库中"""
+        name = await self._get_current_user_name(json.loads(content))
+
+        return text_fletten(f"weibo: [{name[:10]}]")
+
+
 class WeiboSite(Site):
     name = "weibo.com"
     schedule_type = "interval"
     schedule_setting = {"seconds": 3}
+    client_mgr = WeiboClientManager
 
 
 class Weibo(NewMessage):
@@ -78,9 +100,11 @@ class Weibo(NewMessage):
             raise cls.ParseTargetException(prompt="正确格式:\n1. 用户数字UID\n2. https://weibo.com/u/xxxx")
 
     async def get_sub_list(self, target: Target) -> list[RawPost]:
-        client = await self.ctx.get_client()
+        client = await self.ctx.get_client(target)
+        header = {"Referer": f"https://m.weibo.cn/u/{target}", "MWeibo-Pwa": "1", "X-Requested-With": "XMLHttpRequest"}
+        # 获取 cookie 见 https://docs.rsshub.app/zh/deploy/config#%E5%BE%AE%E5%8D%9A
         params = {"containerid": "107603" + target}
-        res = await client.get("https://m.weibo.cn/api/container/getIndex?", params=params, timeout=4.0)
+        res = await client.get("https://m.weibo.cn/api/container/getIndex?", headers=header, params=params, timeout=4.0)
         res_data = json.loads(res.text)
         if not res_data["ok"] and res_data["msg"] != "这里还没有内容":
             raise ApiError(res.request.url)

nonebot_bison/scheduler/manager.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 from nonebot.log import logger
 
 from ..utils import Site
@@ -7,6 +9,7 @@ from ..config.db_model import Target
 from ..types import Target as T_Target
 from ..platform import platform_manager
 from ..plugin_config import plugin_config
+from ..utils.site import CookieClientManager, is_cookie_client_manager
 
 scheduler_dict: dict[type[Site], Scheduler] = {}
 
@@ -41,6 +44,9 @@ async def init_scheduler():
             )
         platform_name_list = _schedule_class_platform_dict[site]
         scheduler_dict[site] = Scheduler(site, schedulable_args, platform_name_list)
+        if is_cookie_client_manager(site.client_mgr):
+            client_mgr = cast(CookieClientManager, scheduler_dict[site].client_mgr)
+            await client_mgr.refresh_client()
     config.register_add_target_hook(handle_insert_new_target)
     config.register_delete_target_hook(handle_delete_target)
 

nonebot_bison/scheduler/scheduler.py

@@ -12,6 +12,7 @@ from ..send import send_msgs
 from ..types import Target, SubUnit
 from ..platform import platform_manager
 from ..utils import Site, ProcessContext
+from ..utils.site import SkipRequestException
 
 
 @dataclass
@@ -107,6 +108,8 @@ class Scheduler:
                     schedulable.platform_name, schedulable.target
                 )
                 to_send = await platform_obj.do_fetch_new_post(SubUnit(schedulable.target, send_userinfo_list))
+        except SkipRequestException as err:
+            logger.debug(f"skip request: {err}")
         except Exception as err:
             records = context.gen_req_records()
             for record in records:

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/__init__.py

@@ -14,6 +14,10 @@ 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 .del_cookie import do_del_cookie
+from .add_cookie_target import do_add_cookie_target
+from .del_cookie_target import do_del_cookie_target
 from .utils import common_platform, admin_permission, gen_handle_cancel, configurable_to_me, set_target_user_info
 
 add_sub_matcher = on_command(
@@ -42,6 +46,46 @@ 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",
+    aliases={"添加Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_add_cookie(add_cookie_matcher)
+
+add_cookie_target_matcher = on_command(
+    "关联cookie",
+    aliases={"关联Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_add_cookie_target(add_cookie_target_matcher)
+
+del_cookie_target_matcher = on_command(
+    "取消关联cookie",
+    aliases={"取消关联Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_del_cookie_target(del_cookie_target_matcher)
+
+del_cookie_matcher = on_command(
+    "删除cookie",
+    aliases={"删除Cookie"},
+    rule=to_me(),
+    permission=SUPERUSER,
+    priority=5,
+    block=True,
+)
+do_del_cookie(del_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, "已取消")
@@ -109,7 +153,11 @@ async def do_dispatch_command(
 
 
 no_permission_matcher = on_command(
-    "添加订阅", rule=configurable_to_me, aliases={"删除订阅", "群管理"}, priority=8, block=True
+    "添加订阅",
+    rule=configurable_to_me,
+    aliases={"删除订阅", "群管理", "管理后台", "添加cookie", "删除cookie", "关联cookie", "取消关联cookie"},
+    priority=8,
+    block=True,
 )
 
 
@@ -125,4 +173,8 @@ __all__ = [
     "del_sub_matcher",
     "group_manage_matcher",
     "no_permission_matcher",
+    "add_cookie_matcher",
+    "add_cookie_target_matcher",
+    "del_cookie_target_matcher",
+    "del_cookie_matcher",
 ]

nonebot_bison/sub_manager/add_cookie.py

@@ -0,0 +1,82 @@
+from typing import cast
+from json import JSONDecodeError
+
+from nonebot.log import logger
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import Arg, ArgPlainText
+from nonebot.adapters.onebot.v11 import MessageEvent
+from nonebot.adapters import Message, MessageTemplate
+
+from ..scheduler import scheduler_dict
+from ..platform import platform_manager
+from ..utils.site import CookieClientManager, is_cookie_client_manager
+from .utils import common_platform, gen_handle_cancel, only_allow_private
+
+
+def do_add_cookie(add_cookie: type[Matcher]):
+    handle_cancel = gen_handle_cancel(add_cookie, "已中止添加cookie")
+
+    @add_cookie.handle()
+    async def init_promote(state: T_State, event: MessageEvent):
+        await only_allow_private(event, add_cookie)
+        state["_prompt"] = (
+            "请输入想要添加 Cookie 的平台，目前支持，请输入冒号左边的名称：\n"
+            + "".join(
+                [
+                    f"{platform_name}: {platform_manager[platform_name].name}\n"
+                    for platform_name in common_platform
+                    if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr)
+                ]
+            )
+            + "要查看全部平台请输入：“全部”\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()
+                    if is_cookie_client_manager(platform_manager[platform_name].site.client_mgr)
+                ]
+            )
+            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(state: T_State):
+        state["_prompt"] = "请输入 Cookie"
+
+    @add_cookie.got("cookie", MessageTemplate("{_prompt}"), [handle_cancel])
+    async def got_cookie(state: T_State, cookie: Message = Arg()):
+        client_mgr = cast(CookieClientManager, scheduler_dict[platform_manager[state["platform"]].site].client_mgr)
+        cookie_text = cookie.extract_plain_text()
+        if not await client_mgr.validate_cookie(cookie_text):
+            await add_cookie.reject(
+                "无效的 Cookie，请检查后重新输入，详情见https://nonebot-bison.netlify.app/usage/cookie.html"
+            )
+        try:
+            cookie_name = await client_mgr.get_cookie_name(cookie_text)
+            state["cookie"] = cookie_text
+            state["cookie_name"] = cookie_name
+        except JSONDecodeError as e:
+            logger.error("获取 Cookie 名称失败" + str(e))
+            await add_cookie.reject(
+                "获取 Cookie 名称失败，请检查后重新输入，详情见https://nonebot-bison.netlify.app/usage/cookie.html"
+            )
+
+    @add_cookie.handle()
+    async def add_cookie_process(state: T_State):
+        client_mgr = cast(CookieClientManager, scheduler_dict[platform_manager[state["platform"]].site].client_mgr)
+        new_cookie = await client_mgr.add_user_cookie(state["cookie"], state["cookie_name"])
+        await add_cookie.finish(
+            f"已添加 Cookie: {new_cookie.cookie_name} 到平台 {state['platform']}"
+            + "\n请使用“关联cookie”为 Cookie 关联订阅"
+        )

nonebot_bison/sub_manager/add_cookie_target.py

@@ -0,0 +1,71 @@
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import ArgPlainText
+from nonebot_plugin_saa import MessageFactory
+from nonebot.adapters.onebot.v11 import MessageEvent
+from nonebot.internal.adapter import MessageTemplate
+
+from ..config import config
+from ..utils import parse_text
+from ..platform import platform_manager
+from .utils import gen_handle_cancel, only_allow_private, generate_sub_list_text
+
+
+def do_add_cookie_target(add_cookie_target_matcher: type[Matcher]):
+    handle_cancel = gen_handle_cancel(add_cookie_target_matcher, "已中止关联 cookie")
+
+    @add_cookie_target_matcher.handle()
+    async def init_promote(state: T_State, event: MessageEvent):
+        await only_allow_private(event, add_cookie_target_matcher)
+        res = await generate_sub_list_text(
+            add_cookie_target_matcher, state, is_index=True, is_show_cookie=True, is_hide_no_cookie_platfrom=True
+        )
+        res += "请输入要关联 cookie 的订阅的序号\n输入'取消'中止"
+        await MessageFactory(await parse_text(res)).send()
+
+    @add_cookie_target_matcher.got("target_idx", parameterless=[handle_cancel])
+    async def got_target_idx(state: T_State, target_idx: str = ArgPlainText()):
+        try:
+            target_idx = int(target_idx)
+            state["target"] = state["sub_table"][target_idx]
+            state["site"] = platform_manager[state["target"]["platform_name"]].site
+        except Exception:
+            await add_cookie_target_matcher.reject("序号错误")
+
+    @add_cookie_target_matcher.handle()
+    async def init_promote_cookie(state: T_State):
+
+        # 获取 site 的所有用户 cookie，再排除掉已经关联的 cookie，剩下的就是可以关联的 cookie
+        cookies = await config.get_cookie(site_name=state["site"].name, is_anonymous=False)
+        associated_cookies = await config.get_cookie(
+            target=state["target"]["target"],
+            site_name=state["site"].name,
+            is_anonymous=False,
+        )
+        associated_cookie_ids = {cookie.id for cookie in associated_cookies}
+        cookies = [cookie for cookie in cookies if cookie.id not in associated_cookie_ids]
+        if not cookies:
+            await add_cookie_target_matcher.finish(
+                "当前平台暂无可关联的 Cookie，请使用“添加cookie”命令添加或检查已关联的 Cookie"
+            )
+        state["cookies"] = cookies
+
+        state["_prompt"] = "请选择一个 Cookie，已关联的 Cookie 不会显示\n" + "\n".join(
+            [f"{idx}. {cookie.cookie_name}" for idx, cookie in enumerate(cookies, 1)]
+        )
+
+    @add_cookie_target_matcher.got("cookie_idx", MessageTemplate("{_prompt}"), [handle_cancel])
+    async def got_cookie_idx(state: T_State, cookie_idx: str = ArgPlainText()):
+        try:
+            cookie_idx = int(cookie_idx)
+            state["cookie"] = state["cookies"][cookie_idx - 1]
+        except Exception:
+            await add_cookie_target_matcher.reject("序号错误")
+
+    @add_cookie_target_matcher.handle()
+    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"]
+        await add_cookie_target_matcher.finish(
+            f"已关联 Cookie: {cookie.cookie_name} " f"到订阅 {state['site'].name} {state['target']['target']}"
+        )

nonebot_bison/sub_manager/del_cookie.py

@@ -0,0 +1,47 @@
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import EventPlainText
+from nonebot_plugin_saa import MessageFactory
+from nonebot.adapters.onebot.v11 import MessageEvent
+
+from ..config import config
+from ..utils import parse_text
+from .utils import gen_handle_cancel, only_allow_private
+
+
+def do_del_cookie(del_cookie: type[Matcher]):
+    handle_cancel = gen_handle_cancel(del_cookie, "删除中止")
+
+    @del_cookie.handle()
+    async def send_list(state: T_State, event: MessageEvent):
+        await only_allow_private(event, del_cookie)
+        cookies = await config.get_cookie(is_anonymous=False)
+        if not cookies:
+            await del_cookie.finish("暂无已添加的 Cookie\n请使用“添加cookie”命令添加")
+        res = "已添加的 Cookie 为：\n"
+        state["cookie_table"] = {}
+        for index, cookie in enumerate(cookies, 1):
+            state["cookie_table"][index] = cookie
+            res += f"{index} {cookie.site_name} {cookie.cookie_name} {len(cookie.targets)}个关联\n"
+            if res[-1] != "\n":
+                res += "\n"
+        res += "请输入要删除的 Cookie 的序号\n输入'取消'中止"
+        await MessageFactory(await parse_text(res)).send()
+
+    @del_cookie.receive(parameterless=[handle_cancel])
+    async def do_del(
+        state: T_State,
+        index_str: str = EventPlainText(),
+    ):
+        try:
+            index = int(index_str)
+            cookie = state["cookie_table"][index]
+            if cookie.targets:
+                await del_cookie.reject("只能删除未关联的 Cookie，请使用“取消关联cookie”命令取消关联")
+            await config.delete_cookie_by_id(cookie.id)
+        except KeyError:
+            await del_cookie.reject("序号错误")
+        except Exception:
+            await del_cookie.reject("删除错误")
+        else:
+            await del_cookie.finish("删除成功")

nonebot_bison/sub_manager/del_cookie_target.py

@@ -0,0 +1,48 @@
+from nonebot.typing import T_State
+from nonebot.matcher import Matcher
+from nonebot.params import EventPlainText
+from nonebot_plugin_saa import MessageFactory
+from nonebot.adapters.onebot.v11 import MessageEvent
+
+from ..config import config
+from ..utils import parse_text
+from .utils import gen_handle_cancel, only_allow_private
+
+
+def do_del_cookie_target(del_cookie_target: type[Matcher]):
+    handle_cancel = gen_handle_cancel(del_cookie_target, "取消关联中止")
+
+    @del_cookie_target.handle()
+    async def send_list(state: T_State, event: MessageEvent):
+        await only_allow_private(event, del_cookie_target)
+        cookie_targets = await config.get_cookie_target()
+        if not cookie_targets:
+            await del_cookie_target.finish("暂无已关联 Cookie\n请使用“添加cookie”命令添加关联")
+        res = "已关联的 Cookie 为：\n"
+        state["cookie_target_table"] = {}
+        for index, cookie_target in enumerate(cookie_targets, 1):
+            friendly_name = cookie_target.cookie.cookie_name
+            state["cookie_target_table"][index] = {
+                "platform_name": cookie_target.target.platform_name,
+                "target": cookie_target.target,
+                "friendly_name": friendly_name,
+                "cookie_target": cookie_target,
+            }
+            res += f"{index} {cookie_target.target.platform_name} {cookie_target.target.target_name} {friendly_name}\n"
+            if res[-1] != "\n":
+                res += "\n"
+        res += "请输入要删除的关联的序号\n输入'取消'中止"
+        await MessageFactory(await parse_text(res)).send()
+
+    @del_cookie_target.receive(parameterless=[handle_cancel])
+    async def do_del(
+        state: T_State,
+        index_str: str = EventPlainText(),
+    ):
+        try:
+            index = int(index_str)
+            await config.delete_cookie_target_by_id(state["cookie_target_table"][index]["cookie_target"].id)
+        except Exception:
+            await del_cookie_target.reject("删除错误")
+        else:
+            await del_cookie_target.finish("删除成功")

nonebot_bison/sub_manager/utils.py

@@ -1,16 +1,21 @@
 import contextlib
 from typing import Annotated
+from itertools import groupby
+from operator import attrgetter
 
 from nonebot.rule import Rule
 from nonebot.adapters import Event
 from nonebot.typing import T_State
 from nonebot.matcher import Matcher
 from nonebot.permission import SUPERUSER
-from nonebot_plugin_saa import extract_target
 from nonebot.params import Depends, EventToMe, EventPlainText
+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 ..utils.site import is_cookie_client_manager
 
 
 def _configurable_to_me(to_me: bool = EventToMe()):
@@ -60,3 +65,67 @@ def admin_permission():
         permission = permission | GROUP_ADMIN | GROUP_OWNER
 
     return permission
+
+
+async def generate_sub_list_text(
+    matcher: type[Matcher],
+    state: T_State,
+    user_info: PlatformTarget | None = None,
+    is_index=False,
+    is_show_cookie=False,
+    is_hide_no_cookie_platfrom=False,
+):
+    """根据配置参数，生产订阅列表文本，同时将订阅信息存入state["sub_table"]"""
+    if user_info:
+        sub_list = await config.list_subscribe(user_info)
+    else:
+        sub_list = await config.list_subs_with_all_info()
+        sub_list = [
+            next(group)
+            for key, group in groupby(sorted(sub_list, key=attrgetter("target_id")), key=attrgetter("target_id"))
+        ]
+    if is_hide_no_cookie_platfrom:
+        sub_list = [
+            sub
+            for sub in sub_list
+            if is_cookie_client_manager(platform_manager.get(sub.target.platform_name).site.client_mgr)
+        ]
+    if not sub_list:
+        await matcher.finish("暂无已订阅账号\n请使用“添加订阅”命令添加订阅")
+    res = "订阅的帐号为：\n"
+    state["sub_table"] = {}
+    for index, sub in enumerate(sub_list, 1):
+        state["sub_table"][index] = {
+            "platform_name": sub.target.platform_name,
+            "target": sub.target.target,
+        }
+        res += f"{index} " if is_index else ""
+        res += f"{sub.target.platform_name} {sub.target.target_name} {sub.target.target}\n"
+        if platform := platform_manager.get(sub.target.platform_name):
+            if platform.categories:
+                res += " [{}]".format(", ".join(platform.categories[Category(x)] for x in sub.categories)) + "\n"
+            if platform.enable_tag:
+                if sub.tags:
+                    res += " {}".format(", ".join(sub.tags)) + "\n"
+            if is_show_cookie:
+                target_cookies = await config.get_cookie(
+                    target=sub.target.target, site_name=platform.site.name, is_anonymous=False
+                )
+                if target_cookies:
+                    res += "  关联的 Cookie：\n"
+                    for cookie in target_cookies:
+                        res += f"  \t{cookie.cookie_name}\n"
+
+        else:
+            res += f" （平台 {sub.target.platform_name} 已失效，请删除此订阅）"
+
+    return res
+
+
+async def only_allow_private(
+    event: Event,
+    matcher: type[Matcher],
+):
+    # if not issubclass(PrivateMessageEvent, event.__class__):
+    if event.message_type != "private":
+        await matcher.finish("请在私聊中使用此命令")

nonebot_bison/utils/context.py

@@ -22,8 +22,9 @@ class ProcessContext:
         async def _log_to_ctx(r: Response):
             self._log_response(r)
 
+        existing_hooks = client.event_hooks["response"]
         hooks = {
-            "response": [_log_to_ctx],
+            "response": [*existing_hooks, _log_to_ctx],
         }
         client.event_hooks = hooks
 

nonebot_bison/utils/site.py

@@ -1,10 +1,18 @@
+import json
 from typing import Literal
+from json import JSONDecodeError
 from abc import ABC, abstractmethod
+from collections.abc import Callable
+from datetime import datetime, timedelta
 
+import httpx
 from httpx import AsyncClient
+from nonebot.log import logger
 
 from ..types import Target
+from ..config import config
 from .http import http_client
+from ..config.db_model import Cookie
 
 
 class ClientManager(ABC):
@@ -35,12 +43,158 @@ class DefaultClientManager(ClientManager):
         pass
 
 
-class Site:
+class SkipRequestException(Exception):
+    """跳过请求异常，如果需要在选择 Cookie 时跳过此次请求，可以抛出此异常"""
+
+    pass
+
+
+class CookieClientManager(ClientManager):
+    _default_cookie_cd = timedelta(seconds=15)
+    _site_name: str = ""
+
+    async def _generate_anonymous_cookie(self) -> Cookie:
+        return Cookie(
+            cookie_name=f"{self._site_name} anonymous",
+            site_name=self._site_name,
+            content="{}",
+            is_universal=True,
+            is_anonymous=True,
+            last_usage=datetime.now(),
+            cd_milliseconds=0,
+            tags="{}",
+            status="",
+        )
+
+    async def _refresh_anonymous_cookie(self):
+        """更新已有的匿名cookie，若不存在则添加"""
+        existing_anonymous_cookies = await config.get_cookie(self._site_name, is_anonymous=True)
+        if existing_anonymous_cookies:
+            for cookie in existing_anonymous_cookies:
+                new_anonymous_cookie = await self._generate_anonymous_cookie()
+                new_anonymous_cookie.id = cookie.id  # 保持原有的id
+                await config.update_cookie(new_anonymous_cookie)
+        else:
+            new_anonymous_cookie = await self._generate_anonymous_cookie()
+            await config.add_cookie(new_anonymous_cookie)
+
+    async def add_user_cookie(self, content: str, cookie_name: str | None = None) -> Cookie:
+        """添加用户 cookie"""
+
+        if not await self.validate_cookie(content):
+            raise ValueError()
+        cookie = Cookie(site_name=self._site_name, content=content)
+        cookie.cookie_name = cookie_name if cookie_name else await self.get_cookie_name(content)
+        cookie.cd = self._default_cookie_cd
+        cookie_id = await config.add_cookie(cookie)
+        return await config.get_cookie_by_id(cookie_id)
+
+    async def get_cookie_name(self, content: str) -> str:
+        """从cookie内容中获取cookie的友好名字，添加cookie时调用，持久化在数据库中"""
+        from . import text_fletten
+
+        return text_fletten(f"{self._site_name} [{content[:10]}]")
+
+    async def validate_cookie(self, content: str) -> bool:
+        """验证 cookie 内容是否有效，添加 cookie 时用，可根据平台的具体情况进行重写"""
+        try:
+            data = json.loads(content)
+            if not isinstance(data, dict):
+                return False
+        except JSONDecodeError:
+            return False
+        return True
+
+    def _generate_hook(self, cookie: Cookie) -> Callable:
+        """hook 函数生成器，用于回写请求状态到数据库"""
+
+        async def _response_hook(resp: httpx.Response):
+            if resp.status_code == 200:
+                logger.trace(f"请求成功: {cookie.id} {resp.request.url}")
+                cookie.status = "success"
+            else:
+                logger.warning(f"请求失败: {cookie.id} {resp.request.url}, 状态码: {resp.status_code}")
+                cookie.status = "failed"
+            cookie.last_usage = datetime.now()
+            await config.update_cookie(cookie)
+
+        return _response_hook
+
+    async def _choose_cookie(self, target: Target | None) -> Cookie:
+        """选择 cookie 的具体算法"""
+        cookies = await config.get_cookie(self._site_name, target)
+        cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
+        cookie = min(cookies, key=lambda x: x.last_usage)
+        return cookie
+
+    async def get_client(self, target: Target | None) -> AsyncClient:
+        """获取 client，根据 target 选择 cookie"""
+        client = http_client()
+        cookie = await self._choose_cookie(target)
+        if cookie.is_universal:
+            logger.trace(f"平台 {self._site_name} 未获取到用户cookie, 使用匿名cookie")
+        else:
+            logger.trace(f"平台 {self._site_name} 获取到用户cookie: {cookie.id}")
+
+        return await self._assemble_client(client, cookie)
+
+    async def _assemble_client(self, client, cookie) -> AsyncClient:
+        """组装 client，可以自定义 cookie 对象装配到 client 中的方式"""
+        cookies = httpx.Cookies()
+        if cookie:
+            cookies.update(json.loads(cookie.content))
+        client.cookies = cookies
+        client.event_hooks = {"response": [self._generate_hook(cookie)]}
+        return client
+
+    @classmethod
+    def from_name(cls, site_name: str) -> type["CookieClientManager"]:
+        """创建一个平台特化的 CookieClientManger"""
+        return type(
+            "CookieClientManager",
+            (CookieClientManager,),
+            {"_site_name": site_name},
+        )
+
+    async def get_client_for_static(self) -> AsyncClient:
+        return http_client()
+
+    async def get_query_name_client(self) -> AsyncClient:
+        return http_client()
+
+    async def refresh_client(self):
+        await self._refresh_anonymous_cookie()
+
+
+def is_cookie_client_manager(manger: type[ClientManager]) -> bool:
+    return issubclass(manger, CookieClientManager)
+
+
+site_manager: dict[str, type["Site"]] = {}
+
+
+class SiteMeta(type):
+    def __new__(cls, name, bases, namespace, **kwargs):
+        return super().__new__(cls, name, bases, namespace)
+
+    def __init__(cls, name, bases, namespace, **kwargs):
+        if kwargs.get("base"):
+            # this is the base class
+            cls._key = kwargs.get("key")
+        elif not kwargs.get("abstract"):
+            # this is the subclass
+            if hasattr(cls, "name"):
+                site_manager[cls.name] = cls
+        super().__init__(name, bases, namespace, **kwargs)
+
+
+class Site(metaclass=SiteMeta):
     schedule_type: Literal["date", "interval", "cron"]
     schedule_setting: dict
     name: str
     client_mgr: type[ClientManager] = DefaultClientManager
     require_browser: bool = False
+    registry: list[type["Site"]]
 
     def __str__(self):
         return f"[{self.name}]-{self.name}-{self.schedule_setting}"