✨ 添加 Cookie 组件 (#633)

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

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)