✨ 添加 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