diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py index 14111c8..da8208c 100644 --- a/nonebot_bison/admin_page/api.py +++ b/nonebot_bison/admin_page/api.py @@ -22,6 +22,7 @@ from ..utils.site import CookieClientManager, site_manager, is_cookie_client_man from ..config import NoSuchUserException, NoSuchTargetException, NoSuchSubscribeException, config from .types import ( Cookie, + Target, TokenResp, GlobalConf, SiteConfig, @@ -211,7 +212,7 @@ async def update_weigth_config(platformName: str, target: str, weight_config: We @router.get("/cookie", dependencies=[Depends(check_is_superuser)]) -async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]: +async def get_cookie(site_name: str | None = None, target: str | None = None) -> list[Cookie]: cookies_in_db = await config.get_cookie(site_name, is_anonymous=False) return [ Cookie( @@ -250,7 +251,12 @@ async def get_cookie_target( cookie_targets = await config.get_cookie_target() # TODO: filter in SQL return [ - x + CookieTarget( + target=Target( + platform_name=x.target.platform_name, target_name=x.target.target_name, target=x.target.target + ), + cookie_id=x.cookie.id, + ) 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) @@ -259,13 +265,13 @@ async def get_cookie_target( @router.post("/cookie_target", dependencies=[Depends(check_is_superuser)]) -async def add_cookie_target(platform_name: str, target: str, cookie_id: int) -> StatusResp: +async def add_cookie_target(platform_name: str, target: T_Target, 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: +async def del_cookie_target(platform_name: str, target: T_Target, cookie_id: int) -> StatusResp: await config.delete_cookie_target(target, platform_name, cookie_id) return StatusResp(ok=True, msg="") diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 7451292..c76078c 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -288,6 +288,8 @@ class DBConfig: 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)) + if not cookie: + raise NoSuchTargetException(f"cookie {cookie_id} not found") return cookie async def add_cookie(self, cookie: Cookie) -> int: @@ -317,6 +319,8 @@ class DBConfig: .outerjoin(CookieTarget) .options(selectinload(Cookie.targets)) ) + if not cookie: + raise NoSuchTargetException(f"cookie {cookie_id} not found") if len(cookie.targets) > 0: raise Exception(f"cookie {cookie.id} in use") await sess.execute(delete(Cookie).where(Cookie.id == cookie_id)) diff --git a/nonebot_bison/config/subs_io/nbesf_model/v3.py b/nonebot_bison/config/subs_io/nbesf_model/v3.py index a2d5e42..0113394 100644 --- a/nonebot_bison/config/subs_io/nbesf_model/v3.py +++ b/nonebot_bison/config/subs_io/nbesf_model/v3.py @@ -8,9 +8,11 @@ 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 ....types import Tag +from ....types import Category from ..utils import NBESFParseErr -from ....types import Tag, Category from .base import NBESFBase, SubReceipt +from ....types import Target as T_Target from ...db_model import Cookie as DBCookie from ...db_config import SubscribeDupException, config @@ -114,7 +116,7 @@ async def magic_cookie_gen(nbesf_data: SubGroup): 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) + await config.add_cookie_target(T_Target(target.target), target.platform_name, cookie_id) except Exception as e: logger.error(f"！添加 Cookie 条目 {repr(cookie)} 失败: {repr(e)}") else: diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 8dcd371..aecb00f 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -123,8 +123,8 @@ class CookieClientManager(ClientManager): 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) + avaliable_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now()) + cookie = min(avaliable_cookies, key=lambda x: x.last_usage) return cookie async def get_client(self, target: Target | None) -> AsyncClient: @@ -183,8 +183,8 @@ class SiteMeta(type): cls._key = kwargs.get("key") elif not kwargs.get("abstract"): # this is the subclass - if hasattr(cls, "name"): - site_manager[cls.name] = cls + if "name" in namespace: + site_manager[namespace["name"]] = cls super().__init__(name, bases, namespace, **kwargs) diff --git a/tests/subs_io/test_cli.py b/tests/subs_io/test_cli.py index d74bac5..61fe271 100644 --- a/tests/subs_io/test_cli.py +++ b/tests/subs_io/test_cli.py @@ -78,7 +78,7 @@ async def test_subs_export(app: App, tmp_path: Path): cookie_name="test cookie", ) ) - await config.add_cookie_target("weibo_id", "weibo", cookie_id) + await config.add_cookie_target(TTarget("weibo_id"), "weibo", cookie_id) assert len(await config.list_subs_with_all_info()) == 3