This commit is contained in:
2025-10-25 22:00:57 +08:00
2 changed files with 79 additions and 99 deletions

View File

@ -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()

View File

@ -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"