import asyncio
import re
import html as html_module
import markdown
try:
import bleach
_HAS_BLEACH = True
except ImportError:
_HAS_BLEACH = False
from nonebot import logger
async def markdown_to_image(markdown_text: str, output_path: str, browser=None):
"""将 Markdown 转换为 HTML 并使用 Puppeteer 截图。"""
page = None
should_close_browser = False
try:
# Convert markdown to HTML. The markdown library handles special chars safely.
# Note: do NOT html.escape() before markdown.markdown() - it breaks markdown syntax.
html_content = markdown.markdown(markdown_text, extensions=["fenced_code", "tables"])
# Sanitize to prevent XSS from malicious AI responses
allowed_tags = [
'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'br', 'hr',
'ul', 'ol', 'li', 'blockquote', 'pre', 'code', 'span',
'table', 'thead', 'tbody', 'tr', 'th', 'td',
'strong', 'em', 'b', 'i', 'u', 'a', 'img', 'div',
]
allowed_attrs = {'a': ['href', 'title'], 'img': ['src', 'alt', 'title']}
if _HAS_BLEACH:
html_content = bleach.clean(html_content, tags=allowed_tags, attributes=allowed_attrs)
else:
logger.warning("[chatai] bleach 未安装，跳过 HTML 净化（请运行 pip install bleach）")
# 使用传入的浏览器实例或创建新的
if browser is None:
from pyppeteer import launch
browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])
should_close_browser = True
page = await browser.newPage()
page.setDefaultNavigationTimeout(15000)
# 设置页面样式，使内容更美观
await page.setContent(f"""
{html_content}
""")
# 等待内容渲染完成
try:
await asyncio.wait_for(page.waitForNavigation({'waitUntil': 'networkidle0'}), timeout=10)
except Exception:
pass # rendering may already be complete
# 获取内容尺寸并设置视口
dimensions = await page.evaluate('''() => {
const container = document.querySelector('.container');
return {
width: container.offsetWidth + 60, // 加上 body 的 padding
height: container.offsetHeight + 60
}
}''')
# 设置视口大小
await page.setViewport({
'width': dimensions['width'],
'height': dimensions['height'],
'deviceScaleFactor': 2 # 提高图片清晰度
})
# 截图，使用透明背景
await page.screenshot({
'path': output_path,
'omitBackground': True, # 透明背景
'clip': {
'x': 0,
'y': 0,
'width': dimensions['width'],
'height': dimensions['height']
}
})
# 关闭页面
await page.close()
# 如果是我们创建的浏览器实例，则关闭它
if should_close_browser:
await browser.close()
except Exception as e:
# 确保资源被释放
if page is not None:
try:
await page.close()
except Exception:
pass
if should_close_browser and browser is not None:
try:
await browser.close()
except Exception:
pass
raise # 重新抛出异常以便上层处理