diff --git a/konabot/common/web_render/__init__.py b/konabot/common/web_render/__init__.py index 06dd688..b9ec029 100644 --- a/konabot/common/web_render/__init__.py +++ b/konabot/common/web_render/__init__.py @@ -1,7 +1,12 @@ import asyncio import queue +from typing import Any, Callable, Coroutine from loguru import logger -from playwright.async_api import async_playwright, Browser, Page, BrowserContext +from playwright.async_api import Page, Playwright, async_playwright, Browser, Page, BrowserContext + + +PageFunction = Callable[[Page], Coroutine[Any, Any, Any]] + class WebRenderer: browser_pool: queue.Queue["WebRendererInstance"] = queue.Queue() @@ -27,7 +32,14 @@ class WebRenderer: return cls.context_pool[id(instance)] @classmethod - async def render(cls, url: str, target: str, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes: + async def render( + cls, + url: str, + target: str, + params: dict = {}, + other_function: PageFunction | None = None, + timeout: int = 30, + ) -> bytes: ''' 访问指定URL并返回截图 @@ -80,21 +92,43 @@ class WebRenderer: del cls.page_pool[page_id] logger.debug(f"Closed and removed persistent page for page_id {page_id}") + + class WebRendererInstance: def __init__(self): - self.playwright = None - self.browser: Browser = None - self.lock: asyncio.Lock = None + self._playwright: Playwright | None = None + self._browser: Browser | None = None + self.lock = asyncio.Lock() + + @property + def playwright(self) -> Playwright: + assert self._playwright is not None + return self._playwright + + @property + def browser(self) -> Browser: + assert self._browser is not None + return self._browser + + async def init(self): + self._playwright = await async_playwright().start() + self._browser = await self.playwright.chromium.launch(headless=True) @classmethod async def create(cls) -> "WebRendererInstance": instance = cls() - instance.playwright = await async_playwright().start() - instance.browser = await instance.playwright.chromium.launch(headless=True) - instance.lock = asyncio.Lock() + await instance.init() return instance - async def render(self, url: str, target: str, index: int = 0, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes: + async def render( + self, + url: str, + target: str, + index: int = 0, + params: dict = {}, + other_function: PageFunction | None = None, + timeout: int = 30 + ) -> bytes: ''' 访问指定URL并返回截图 @@ -142,7 +176,30 @@ class WebRendererInstance: logger.debug(f"Screenshot taken successfully") return screenshot + logger.debug(f"Navigating to {url} with timeout {timeout}") + try: + url_with_params = url + ("?" + "&".join(f"{k}={v}" for k, v in params.items()) if params else "") + await page.goto(url_with_params, timeout=timeout * 1000, wait_until="load") + logger.debug("Page loaded successfully") + # 等待目标元素出现 + await page.wait_for_selector(target, timeout=timeout * 1000) + logger.debug(f"Target element '{target}' found, taking screenshot") + if other_function: + await other_function(page) + elements = await page.query_selector_all(target) + if not elements: + raise Exception(f"Target element '{target}' not found on the page.") + if index >= len(elements): + raise Exception(f"Index {index} out of range for elements matching '{target}'.") + element = elements[index] + screenshot = await element.screenshot() + logger.debug("Screenshot taken successfully") + return screenshot + finally: + await page.close() + await context.close() async def close(self): await self.browser.close() await self.playwright.stop() + diff --git a/konabot/plugins/simple_notify/__init__.py b/konabot/plugins/simple_notify/__init__.py index 9960c0c..431ee7e 100644 --- a/konabot/plugins/simple_notify/__init__.py +++ b/konabot/plugins/simple_notify/__init__.py @@ -1,23 +1,19 @@ import aiohttp import asyncio as asynkio -import datetime from math import ceil from pathlib import Path -from typing import Any, Literal +from typing import Any import nanoid import nonebot import ptimeparse from loguru import logger from nonebot import get_plugin_config, on_message -from nonebot.adapters import Bot, Event -from nonebot.adapters.onebot.v11 import Bot as OBBot -from nonebot.adapters.console import Bot as CBot -from nonebot.adapters.discord import Bot as DCBot +from nonebot.adapters import Event from nonebot_plugin_alconna import Alconna, Args, Subcommand, UniMessage, UniMsg, on_alconna from pydantic import BaseModel -from konabot.common.longtask import DepLongTaskTarget, LongTask, LongTaskTarget, create_longtask, handle_long_task, longtask_data +from konabot.common.longtask import DepLongTaskTarget, LongTask, create_longtask, handle_long_task, longtask_data evt = on_message() @@ -32,27 +28,8 @@ PAGE_SIZE = 6 FMT_STRING = "%Y年%m月%d日 %H:%M:%S" -class NotifyMessage(BaseModel): - message: str - - -class Notify(BaseModel): - platform: Literal["console", "qq", "discord"] - - target: str - "需要接受通知的个体" - - target_env: str | None - "在哪里进行通知，如果是 None 代表私聊通知" - - notify_time: datetime.datetime - notify_msg: str - - class NotifyConfigFile(BaseModel): version: int = 2 - notifies: list[Notify] = [] - unsent: list[Notify] = [] notify_channels: dict[str, str] = {} @@ -78,36 +55,6 @@ async def send_notify_to_ntfy_instance(msg: str, channel: str): logger.info(f"访问 {url} 的结果是 {response.status}") -def _get_bot_of(_type: type[Bot]): - for bot in nonebot.get_bots().values(): - if isinstance(bot, _type): - return bot.self_id - return "" - - -def get_target_from_notify(notify: Notify) -> LongTaskTarget: - if notify.platform == "console": - return LongTaskTarget( - platform="console", - self_id=_get_bot_of(CBot), - channel_id=notify.target_env or "", - target_id=notify.target, - ) - if notify.platform == "discord": - return LongTaskTarget( - platform="discord", - self_id=_get_bot_of(DCBot), - channel_id=notify.target_env or "", - target_id=notify.target, - ) - return LongTaskTarget( - platform="qq", - self_id=_get_bot_of(OBBot), - channel_id=notify.target_env or "", - target_id=notify.target, - ) - - def load_notify_config() -> NotifyConfigFile: if not DATA_FILE_PATH.exists(): return NotifyConfigFile() @@ -160,40 +107,6 @@ async def _(msg: UniMsg, mEvt: Event, target: DepLongTaskTarget): driver = nonebot.get_driver() -NOTIFIED_FLAG = { - "task_added": False, -} - - -@driver.on_bot_connect -async def _(): - if NOTIFIED_FLAG["task_added"]: - return - - NOTIFIED_FLAG["task_added"] = True - - DELTA = 2 - logger.info(f"第一次探测到 Bot 连接，等待 {DELTA} 秒后开始通知") - await asynkio.sleep(DELTA) - - await DATA_FILE_LOCK.acquire() - - cfg = load_notify_config() - if cfg.version == 1: - logger.info("将配置文件的版本升级为 2") - cfg.version = 2 - else: - for notify in [*cfg.notifies]: - await create_longtask( - handler=LONG_TASK_NAME, - data={ "message": notify.notify_msg }, - target=get_target_from_notify(notify), - deadline=notify.notify_time, - ) - cfg.notifies = [] - save_notify_config(cfg) - DATA_FILE_LOCK.release() - @handle_long_task("TASK_SIMPLE_NOTIFY") async def _(task: LongTask): @@ -284,7 +197,17 @@ cmd_notify_channel = on_alconna(Alconna( @cmd_notify_channel.assign("$main") async def _(target: DepLongTaskTarget): + async with DATA_FILE_LOCK: + data = load_notify_config() + target_channel = data.notify_channels.get(target.target_id) + + if target_channel is None: + channel_msg = "目前还没有配置 ntfy 地址" + else: + channel_msg = f"配置的 ntfy Channel 为：{target_channel}\n\n服务器地址：{config.plugin_notify_base_url}" + await target.send_message(UniMessage.text( + f"{channel_msg}\n\n" "配置 ntfy 通知：\n\n" "- ntfy 创建: 启用 ntfy 通知，并为你随机生成一个通知渠道\n" "- ntfy 删除：禁用 ntfy 通知\n"