diff --git a/assets/img/other/boom.jpg b/assets/img/other/boom.jpg new file mode 100644 index 0000000..dfd333b Binary files /dev/null and b/assets/img/other/boom.jpg differ diff --git a/konabot/common/web_render/__init__.py b/konabot/common/web_render/__init__.py index 77f000e..06dd688 100644 --- a/konabot/common/web_render/__init__.py +++ b/konabot/common/web_render/__init__.py @@ -1,10 +1,30 @@ import asyncio import queue from loguru import logger -from playwright.async_api import async_playwright, Browser +from playwright.async_api import async_playwright, Browser, Page, BrowserContext class WebRenderer: browser_pool: queue.Queue["WebRendererInstance"] = queue.Queue() + context_pool: dict[int, BrowserContext] = {} # 长期挂载的浏览器上下文池 + page_pool: dict[str, Page] = {} # 长期挂载的页面池 + + @classmethod + async def get_browser_instance(cls) -> "WebRendererInstance": + if cls.browser_pool.empty(): + instance = await WebRendererInstance.create() + cls.browser_pool.put(instance) + instance = cls.browser_pool.get() + cls.browser_pool.put(instance) + return instance + + @classmethod + async def get_browser_context(cls) -> BrowserContext: + instance = await cls.get_browser_instance() + if id(instance) not in cls.context_pool: + context = await instance.browser.new_context() + cls.context_pool[id(instance)] = context + logger.debug(f"Created new persistent browser context for WebRendererInstance {id(instance)}") + 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: @@ -19,14 +39,46 @@ class WebRenderer: :return: 截图的字节数据 ''' - logger.debug(f"Requesting render for {url} targeting {target} with timeout {timeout}") - if cls.browser_pool.empty(): - instance = await WebRendererInstance.create() - cls.browser_pool.put(instance) - instance = cls.browser_pool.get() - cls.browser_pool.put(instance) + instance = await cls.get_browser_instance() logger.debug(f"Using WebRendererInstance {id(instance)} to render {url} targeting {target}") return await instance.render(url, target, params=params, other_function=other_function, timeout=timeout) + + @classmethod + async def render_persistent_page(cls, page_id: str, url: str, target: str, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes: + ''' + 使用长期挂载的页面访问指定URL并返回截图 + + :param page_id: 页面唯一标识符 + :param url: 目标URL + :param target: 渲染目标,如 ".box"、"#main" 等CSS选择器 + :param timeout: 页面加载超时时间,单位秒 + :param params: URL键值对参数 + :param other_function: 其他自定义操作函数,接受page参数 + :return: 截图的字节数据 + + ''' + logger.debug(f"Requesting persistent render for page_id {page_id} at {url} targeting {target} with timeout {timeout}") + instance = await cls.get_browser_instance() + if page_id not in cls.page_pool: + context = await cls.get_browser_context() + page = await context.new_page() + cls.page_pool[page_id] = page + logger.debug(f"Created new persistent page for page_id {page_id} using WebRendererInstance {id(instance)}") + page = cls.page_pool[page_id] + return await instance.render_with_page(page, url, target, params=params, other_function=other_function, timeout=timeout) + + @classmethod + async def close_persistent_page(cls, page_id: str) -> None: + ''' + 关闭并移除长期挂载的页面 + + :param page_id: 页面唯一标识符 + ''' + if page_id in cls.page_pool: + page = cls.page_pool[page_id] + await page.close() + del cls.page_pool[page_id] + logger.debug(f"Closed and removed persistent page for page_id {page_id}") class WebRendererInstance: def __init__(self): @@ -58,28 +110,38 @@ class WebRendererInstance: async with self.lock: context = await self.browser.new_context() page = await context.new_page() - 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(f"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(f"Screenshot taken successfully") - return screenshot - finally: - await page.close() - await context.close() + screenshot = await self.inner_render(page, url, target, index, params, other_function, timeout) + await page.close() + await context.close() + return screenshot + + async def render_with_page(self, page: Page, url: str, target: str, index: int = 0, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes: + async with self.lock: + screenshot = await self.inner_render(page, url, target, index, params, other_function, timeout) + return screenshot + + async def inner_render(self, page: Page, url: str, target: str, index: int = 0, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes: + logger.debug(f"Navigating to {url} with timeout {timeout}") + 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(f"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: + logger.error(f"Target element '{target}' not found on the page.") + return None + if index >= len(elements): + logger.error(f"Index {index} out of range for elements matching '{target}'") + return None + element = elements[index] + screenshot = await element.screenshot() + logger.debug(f"Screenshot taken successfully") + return screenshot + async def close(self): await self.browser.close() diff --git a/konabot/plugins/air_conditioner/__init__.py b/konabot/plugins/air_conditioner/__init__.py new file mode 100644 index 0000000..12f8e29 --- /dev/null +++ b/konabot/plugins/air_conditioner/__init__.py @@ -0,0 +1,91 @@ +from typing import Optional +from nonebot_plugin_alconna import Alconna, Args, UniMessage, UniMsg, on_alconna +from konabot.common.longtask import DepLongTaskTarget +from konabot.common.path import ASSETS_PATH +from konabot.common.web_render import WebRenderer +from nonebot.adapters import Event as BaseEvent +from nonebot.adapters.discord.event import MessageEvent as DiscordMessageEvent +from playwright.async_api import Page + +async def open_handle(page: Page) -> None: + ''' + 开空调 + ''' + # 找到 id 为 power 的开关按钮元素 + power_button = await page.query_selector("#power") + if power_button: + # 点击按钮打开空调 + await power_button.click(force=True) + +async def up_handle(page: Page) -> None: + ''' + 升温 + ''' + # 找到 id 为 add 的按钮元素 + add_button = await page.query_selector("#add") + if add_button: + # 点击按钮升温,无需检测是否稳定 + await add_button.click(force=True) + +async def down_handle(page: Page) -> None: + ''' + 降温 + ''' + # 找到 id 为 minus 的按钮元素 + minus_button = await page.query_selector("#minus") + if minus_button: + # 点击按钮降温 + await minus_button.click(force=True) + +def get_user_info(event: BaseEvent): + if isinstance(event, DiscordMessageEvent): + user_id = str(event.author.id) + user_name = str(event.author.name) + else: + user_id = str(event.get_user_id()) + user_name = str(event.get_user_id()) + return user_id, user_name + +evt = on_alconna( + Alconna( + f"群空调", + Args["condition", str] + ), + use_cmd_start=True, + use_cmd_sep=False, + skip_for_unmatch=True, +) +@evt.handle() +async def _(msg: UniMsg, event: BaseEvent, target: DepLongTaskTarget, condition: Optional[str] = ""): + identify_code = f"air_conditioner_{target.channel_id}" + function_handle = None + match condition: + case "开空调" | "打开空调" | "启动空调" | "关闭空调" | "关空调" | "开关空调": + function_handle = open_handle + case "升温" | "调高温度" | "加温" | "加" | "调高" | "提高" | "加一度": + function_handle = up_handle + case "降温" | "调低温度" | "减温" | "减" | "调低" | "降低" | "减一度": + function_handle = down_handle + case "炸空调": + await WebRenderer.close_persistent_page(identify_code) + # 读取 boom 图片 + with open(ASSETS_PATH / "img" / "other" / "boom.jpg", "rb") as f: + boom_image = f.read() + await evt.send(await UniMessage().image(raw=boom_image).export()) + user_id, _ = get_user_info(event) + await evt.send(await UniMessage().at(user_id).text("空调被你炸毁了!我们重新装了一台!").export()) + return + case _: + return + + screenshot = await WebRenderer.render_persistent_page( + page_id=identify_code, + url="https://toolwa.com/ac/", + target="#kt", + other_function=lambda page: function_handle(page) if function_handle else None, + timeout=30 + ) + + await evt.send( + await UniMessage().image(raw=screenshot).export() + ) \ No newline at end of file