from typing import ParamSpec from functools import partial from datetime import timedelta from collections import defaultdict from httpx import AsyncClient from nonebot import logger, require from rapidfuzz import fuzz, process from nonebot_bison.post import Post from nonebot_bison.plugin_config import plugin_config from nonebot_bison.types import Target, RawPost, Category from nonebot_bison.utils import Site, ClientManager, capture_html from ..platform import NewMessage from .utils import process_response from .const import COMB_ID_URL, COOKIES_URL, COOKIE_ID_URL from .exception import CeobeSnapshotSkip, CeobeSnapshotFailed from .cache import CeobeCache, CeobeClient, CeobeDataSourceCache from .models import CeobeImage, CeobeCookie, CeobeTextPic, CombIdResponse, CookiesResponse, CookieIdResponse P = ParamSpec("P") class CeobeCanteenClientManager(ClientManager): _client: AsyncClient def __init__(self): self._client = CeobeClient( headers={ "User-Agent": "MountainDash/Nonebot-Bison", } ) async def get_client(self, target: Target | None) -> AsyncClient: return self._client async def get_client_for_static(self) -> AsyncClient: return self._client async def get_query_name_client(self) -> AsyncClient: return self._client async def refresh_client(self): raise NotImplementedError("refresh_client is not implemented") class CeobeCanteenSite(Site): name = "ceobe_canteen" schedule_type = "interval" # lwt の 推荐间隔 schedule_setting = {"seconds": 15} client_mgr = CeobeCanteenClientManager class CeobeCanteen(NewMessage): enable_tag: bool = False platform_name: str = "ceobecanteen" name: str = "小刻食堂" enabled: bool = True is_common: bool = False site = CeobeCanteenSite has_target: bool = True use_batch: bool = True default_theme: str = "ceobecanteen" categories: dict[Category, str] = {1: "普通", 2: "转发"} data_source_cache = CeobeDataSourceCache() comb_id = CeobeCache(timedelta(hours=12)) cookie_id = CeobeCache(timedelta(hours=1)) cookies = CeobeCache(timedelta(hours=1)) async def get_comb_id(self, target_uuids: list[str]): """获取数据源的组合id""" payload = {"datasource_push": target_uuids} logger.trace(payload) client = await self.ctx.get_client() resp = await client.post( COMB_ID_URL, json=payload, ) comb_id = process_response(resp, CombIdResponse).data["datasource_comb_id"] logger.trace(f"get comb_id: {comb_id}") return comb_id async def get_comb_id_of_all(self): """获取 "全部数据源" 的组合id，获取到的comb_id会缓存12小时""" logger.trace("no comb_id, request") target_uuids = (await self.data_source_cache.get_all()).keys() comb_id = await self.get_comb_id(list(target_uuids)) logger.trace(f"use comb_id: {comb_id}") return comb_id async def get_latest_cookie_id(self, comb_id: str): """根据comb_id获取最新cookie_id""" client = await self.ctx.get_client() resp = await client.get(f"{COOKIE_ID_URL}/{comb_id}") cookie_id = process_response(resp, CookieIdResponse).cookie_id logger.trace(f"get cookie_id: {cookie_id}") return cookie_id async def get_cookies(self, cookie_id: str, comb_id: str | None = None): """根据cookie_id、comb_id组合获取cookies""" client = await self.ctx.get_client() parmas = { "datasource_comb_id": comb_id, "cookie_id": cookie_id, } logger.trace(f"will reuquest: {parmas}") resp = await client.get(COOKIES_URL, params=parmas) return process_response(resp, CookiesResponse).data.cookies async def fetch_ceobe_cookies(self) -> list[CeobeCookie]: if not self.comb_id: self.comb_id = await self.get_comb_id_of_all() latest_cookie_id = await self.get_latest_cookie_id(self.comb_id) if not latest_cookie_id: return [] if latest_cookie_id != self.cookie_id: self.cookie_id = latest_cookie_id self.cookies = await self.get_cookies(latest_cookie_id, self.comb_id) return self.cookies or [] async def batch_get_sub_list(self, targets: list[Target]) -> list[list[CeobeCookie]]: cookies = await self.fetch_ceobe_cookies() dispatched_cookies: defaultdict[Target, list[CeobeCookie]] = defaultdict(list) for cookie in cookies: if ceobe_target := await self.data_source_cache.get_by_source(cookie.source): dispatched_cookies[Target(ceobe_target.unique_id)].append(cookie) return [dispatched_cookies[target] for target in targets] @classmethod async def get_target_name(cls, _, uuid_target: Target) -> str: ceobe_target = await cls.data_source_cache.get_by_unique_id(uuid_target) return ceobe_target.nickname if ceobe_target else "UNKNOWN" @classmethod async def parse_target(cls, nickname: str) -> Target: ceobe_target = await cls.data_source_cache.get_by_nickname(nickname) if not ceobe_target: all_targets_name = [target.nickname for target in (await cls.data_source_cache.get_all()).values()] matched_targets_name = process.extract(nickname, all_targets_name, scorer=fuzz.token_sort_ratio, limit=3) logger.debug(f"possible targets: {matched_targets_name}") raise cls.ParseTargetException( prompt="未能匹配到对应的小刻食堂数据源，可能的选择有: \n" + "\n".join([name for name, *_ in matched_targets_name]) + f"\n\n请检查原输入是否正确: {nickname}" ) return Target(ceobe_target.unique_id) def get_tags(self, _: RawPost) -> None: return def get_category(self, post: CeobeCookie) -> Category: if post.item.is_retweeted: return Category(2) return Category(1) def get_id(self, post: CeobeCookie) -> str: return post.item.id def get_date(self, post: CeobeCookie) -> int: return post.timestamp.fetcher async def parse(self, raw_post: CeobeCookie) -> Post: target = await self.data_source_cache.get_by_source(raw_post.source) assert target, "target not found" content, pics = await self.take_snapshot(raw_post) timestamp = raw_post.timestamp.platform or raw_post.timestamp.fetcher if timestamp: timestamp /= 1000 # 从毫秒级转换到秒级 retweet: Post | None = None if raw_post.item.is_retweeted and raw_post.item.retweeted: raw_retweet_pics = raw_post.item.retweeted.images or [] retweet_pics = await self.parse_retweet_images(raw_retweet_pics, raw_post.source.type) retweet = Post( self, nickname=raw_post.item.retweeted.author_name, avatar=raw_post.item.retweeted.author_avatar, images=list(retweet_pics), content=raw_post.item.retweeted.text, ) return Post( self, content, url=raw_post.item.url, nickname=raw_post.datasource, images=list(pics), timestamp=timestamp, avatar=target.avatar, description=target.platform, repost=retweet, ) async def snapshot_official_website(self, url: str) -> bytes: """截取小刻官网的截图""" require("nonebot_plugin_htmlrender") from nonebot_plugin_htmlrender import get_new_page logger.debug(f"snapshot official website url: {url}") # /html/body/div[1]/div[1]/div/div[1]/div[1]/div snapshot_selector = "html > body > div:nth-child(1) > div:nth-child(1) > div > div:nth-child(1) > div:nth-child(1) > div" # noqa: E501 # /html/body/div[1]/div[1]/div/div[1]/div[1]/div/div[4]/div/div/div calculate_selector = "html > body > div:nth-child(1) > div:nth-child(1) > div > div:nth-child(1) > div:nth-child(1) > div > div:nth-child(4) > div > div > div" # noqa: E501 viewport = {"width": 1024, "height": 19990} try: async with get_new_page(viewport=viewport) as page: await page.goto(url, wait_until="networkidle") element_width = await page.evaluate( "(selector) => document.querySelector(selector).offsetWidth", calculate_selector ) logger.debug(f"element width: {element_width}") element_height = await page.evaluate( "(selector) => document.querySelector(selector).offsetHeight", calculate_selector ) logger.debug(f"element height: {element_height}") element_height += 1000 await page.set_viewport_size({"width": 1024, "height": element_height}) element = await page.locator(snapshot_selector).element_handle() # add padding to make the screenshot more beautiful await element.evaluate("(element) => {element.style.padding = '20px';}", element) pic_data = await element.screenshot( type="png", ) except Exception as e: raise CeobeSnapshotFailed("渲染错误") from e else: return pic_data async def snapshot_bulletin_list(self, url: str) -> bytes: """截取小刻公告列表的截图""" selector = "body > div.main > div.container" viewport = {"width": 1024, "height": 19990} try: pic_data = await capture_html( url, selector, timeout=30000, wait_until="networkidle", viewport=viewport, ) assert pic_data except Exception: raise CeobeSnapshotFailed("渲染错误") else: return pic_data async def take_snapshot( self, raw_post: CeobeCookie, ) -> CeobeTextPic: """判断数据源类型，判断是否需要截图""" match raw_post.source.type: case "arknights-website:official-website": async def owss(url: str) -> CeobeTextPic: return CeobeTextPic(text="", pics=[await self.snapshot_official_website(url)]) snapshot_func = partial(owss, raw_post.item.url) case "arknights-game:bulletin-list" if raw_post.item.display_type != 2: async def blss(url: str) -> CeobeTextPic: return CeobeTextPic(text="", pics=[await self.snapshot_bulletin_list(url)]) snapshot_func = partial(blss, raw_post.item.url) case _: async def npss() -> CeobeTextPic: raise CeobeSnapshotSkip("无需截图的数据源") snapshot_func = partial(npss) raw_pics = raw_post.default_cookie.images or [] try: if not plugin_config.bison_use_browser: raise CeobeSnapshotSkip("未启用浏览器") res = await snapshot_func() except CeobeSnapshotSkip as e: logger.info(f"skip snapshot: {e}") pics = await self.parse_retweet_images(raw_pics, raw_post.source.type) res = CeobeTextPic(text=raw_post.default_cookie.text, pics=list(pics)) except CeobeSnapshotFailed: logger.exception("snapshot failed") pics = await self.parse_retweet_images(raw_pics, raw_post.source.type) res = CeobeTextPic(text=raw_post.default_cookie.text, pics=list(pics)) return res async def parse_retweet_images(self, images: list[CeobeImage], source_type: str) -> list[bytes] | list[str]: if source_type.startswith("weibo"): retweet_pics = await self.download_weibo_image([image.origin_url for image in images]) else: retweet_pics = [image.origin_url for image in images] return retweet_pics async def download_weibo_image(self, image_urls: list[str]) -> list[bytes]: headers = {"referer": "https://weibo.cn/"} pics = [] async with CeobeClient(headers=headers) as client: for url in image_urls: resp = await client.get(url) resp.raise_for_status() pics.append(resp.content) return pics