fb5fae1524
- Added new entries to .gitignore for database files. - Updated go.mod and go.sum to include new indirect dependencies for database and ORM support. - Refactored event handling to improve message reply functionality in the protocol. - Enhanced the dispatcher to allow for better event processing and logging. - Removed outdated plugin documentation and unnecessary files to streamline the codebase. - Improved welcome message formatting and screenshot options for better user experience.
352 lines
10 KiB
Go
352 lines
10 KiB
Go
package utils
|
||
|
||
import (
|
||
"bytes"
|
||
"cellbot/internal/protocol"
|
||
"context"
|
||
"encoding/base64"
|
||
"fmt"
|
||
"html/template"
|
||
"os"
|
||
"path/filepath"
|
||
"time"
|
||
|
||
"github.com/chromedp/chromedp"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// ScreenshotOptions 截图选项
|
||
type ScreenshotOptions struct {
|
||
Width int // 视口宽度（像素）
|
||
Height int // 视口高度（像素）
|
||
Timeout time.Duration // 超时时间
|
||
WaitTime time.Duration // 等待时间（页面加载后等待）
|
||
FullPage bool // 是否截取整个页面
|
||
Quality int // 图片质量（0-100，仅PNG格式）
|
||
Format string // 图片格式：png, jpeg
|
||
Logger *zap.Logger // 日志记录器
|
||
}
|
||
|
||
// DefaultScreenshotOptions 默认截图选项
|
||
func DefaultScreenshotOptions() *ScreenshotOptions {
|
||
return &ScreenshotOptions{
|
||
Width: 1920,
|
||
Height: 1080,
|
||
Timeout: 30 * time.Second,
|
||
WaitTime: 1 * time.Second,
|
||
FullPage: false,
|
||
Quality: 90,
|
||
Format: "png",
|
||
Logger: zap.NewNop(),
|
||
}
|
||
}
|
||
|
||
// ScreenshotURL 对指定URL进行截图并返回base64编码
|
||
func ScreenshotURL(ctx context.Context, url string, opts *ScreenshotOptions) (string, error) {
|
||
if opts == nil {
|
||
opts = DefaultScreenshotOptions()
|
||
}
|
||
|
||
// 创建上下文，添加优化选项
|
||
allocCtx, cancel := chromedp.NewExecAllocator(ctx,
|
||
chromedp.NoSandbox,
|
||
chromedp.NoFirstRun,
|
||
chromedp.NoDefaultBrowserCheck,
|
||
chromedp.Headless,
|
||
chromedp.DisableGPU,
|
||
)
|
||
defer cancel()
|
||
|
||
ctx, cancel = chromedp.NewContext(allocCtx, chromedp.WithLogf(func(format string, v ...interface{}) {
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug(fmt.Sprintf(format, v...))
|
||
}
|
||
}))
|
||
defer cancel()
|
||
|
||
// 设置超时
|
||
ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
|
||
defer cancel()
|
||
|
||
var buf []byte
|
||
|
||
// 执行截图任务
|
||
var err error
|
||
if opts.FullPage {
|
||
err = chromedp.Run(ctx,
|
||
chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height)), // 设置视口大小
|
||
chromedp.Navigate(url),
|
||
chromedp.WaitReady("body", chromedp.ByQuery), // 使用 WaitReady 等待页面完全加载
|
||
chromedp.Sleep(opts.WaitTime),
|
||
chromedp.FullScreenshot(&buf, opts.Quality),
|
||
)
|
||
} else {
|
||
err = chromedp.Run(ctx,
|
||
chromedp.Navigate(url),
|
||
chromedp.WaitReady("body", chromedp.ByQuery), // 使用 WaitReady 等待页面完全加载
|
||
chromedp.Sleep(opts.WaitTime),
|
||
chromedp.CaptureScreenshot(&buf),
|
||
)
|
||
}
|
||
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to capture screenshot: %w", err)
|
||
}
|
||
|
||
// 转换为base64
|
||
base64Str := base64.StdEncoding.EncodeToString(buf)
|
||
return base64Str, nil
|
||
}
|
||
|
||
// ScreenshotHTML 对HTML内容进行截图并返回base64编码
|
||
func ScreenshotHTML(ctx context.Context, htmlContent string, opts *ScreenshotOptions) (string, error) {
|
||
if opts == nil {
|
||
opts = DefaultScreenshotOptions()
|
||
}
|
||
|
||
if opts.Logger != nil {
|
||
opts.Logger.Info("Starting HTML screenshot",
|
||
zap.Int("html_length", len(htmlContent)),
|
||
zap.Int("width", opts.Width),
|
||
zap.Int("height", opts.Height),
|
||
zap.Duration("timeout", opts.Timeout))
|
||
}
|
||
|
||
// 创建上下文
|
||
allocCtx, cancel := chromedp.NewExecAllocator(ctx,
|
||
chromedp.NoSandbox,
|
||
chromedp.NoFirstRun,
|
||
chromedp.NoDefaultBrowserCheck,
|
||
chromedp.Headless,
|
||
chromedp.DisableGPU,
|
||
)
|
||
defer cancel()
|
||
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug("Chrome allocator created")
|
||
}
|
||
|
||
ctx, cancel = chromedp.NewContext(allocCtx, chromedp.WithLogf(func(format string, v ...interface{}) {
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug(fmt.Sprintf(format, v...))
|
||
}
|
||
}))
|
||
defer cancel()
|
||
|
||
// 设置超时
|
||
ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
|
||
defer cancel()
|
||
|
||
var buf []byte
|
||
|
||
// 使用 base64 编码的 data URL，避免 URL 编码导致的 + 号问题
|
||
htmlBytes := []byte(htmlContent)
|
||
htmlBase64 := base64.StdEncoding.EncodeToString(htmlBytes)
|
||
dataURL := fmt.Sprintf("data:text/html;charset=utf-8;base64,%s", htmlBase64)
|
||
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug("Navigating to base64 data URL",
|
||
zap.Int("html_length", len(htmlContent)),
|
||
zap.Int("base64_length", len(htmlBase64)))
|
||
}
|
||
|
||
// 执行截图任务
|
||
var err error
|
||
if opts.FullPage {
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug("Taking full page screenshot")
|
||
}
|
||
err = chromedp.Run(ctx,
|
||
chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height)), // 设置视口大小
|
||
chromedp.Navigate(dataURL),
|
||
chromedp.WaitReady("body", chromedp.ByQuery),
|
||
chromedp.Sleep(opts.WaitTime),
|
||
chromedp.FullScreenshot(&buf, opts.Quality),
|
||
)
|
||
} else {
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug("Taking viewport screenshot")
|
||
}
|
||
err = chromedp.Run(ctx,
|
||
chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height)), // 设置视口大小
|
||
chromedp.Navigate(dataURL),
|
||
chromedp.WaitReady("body", chromedp.ByQuery),
|
||
chromedp.Sleep(opts.WaitTime),
|
||
chromedp.CaptureScreenshot(&buf),
|
||
)
|
||
}
|
||
|
||
if err != nil {
|
||
if opts.Logger != nil {
|
||
opts.Logger.Error("Failed to capture screenshot", zap.Error(err))
|
||
}
|
||
return "", fmt.Errorf("failed to capture screenshot: %w", err)
|
||
}
|
||
|
||
if opts.Logger != nil {
|
||
opts.Logger.Info("Screenshot captured successfully",
|
||
zap.Int("image_size", len(buf)))
|
||
}
|
||
|
||
// 转换为base64
|
||
base64Str := base64.StdEncoding.EncodeToString(buf)
|
||
if opts.Logger != nil {
|
||
opts.Logger.Info("Screenshot converted to base64",
|
||
zap.Int("base64_length", len(base64Str)))
|
||
}
|
||
return base64Str, nil
|
||
}
|
||
|
||
// ScreenshotElement 对页面中的特定元素进行截图
|
||
func ScreenshotElement(ctx context.Context, url string, selector string, opts *ScreenshotOptions) (string, error) {
|
||
if opts == nil {
|
||
opts = DefaultScreenshotOptions()
|
||
}
|
||
|
||
// 创建上下文
|
||
allocCtx, cancel := chromedp.NewExecAllocator(ctx,
|
||
chromedp.NoSandbox,
|
||
chromedp.NoFirstRun,
|
||
chromedp.NoDefaultBrowserCheck,
|
||
chromedp.Headless,
|
||
chromedp.DisableGPU,
|
||
)
|
||
defer cancel()
|
||
|
||
ctx, cancel = chromedp.NewContext(allocCtx, chromedp.WithLogf(func(format string, v ...interface{}) {
|
||
if opts.Logger != nil {
|
||
opts.Logger.Debug(fmt.Sprintf(format, v...))
|
||
}
|
||
}))
|
||
defer cancel()
|
||
|
||
// 设置超时
|
||
ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
|
||
defer cancel()
|
||
|
||
var buf []byte
|
||
|
||
// 执行截图任务
|
||
err := chromedp.Run(ctx,
|
||
chromedp.Navigate(url),
|
||
chromedp.WaitVisible(selector, chromedp.ByQuery),
|
||
chromedp.Sleep(opts.WaitTime),
|
||
chromedp.Screenshot(selector, &buf, chromedp.NodeVisible),
|
||
)
|
||
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to capture element screenshot: %w", err)
|
||
}
|
||
|
||
// 转换为base64
|
||
base64Str := base64.StdEncoding.EncodeToString(buf)
|
||
return base64Str, nil
|
||
}
|
||
|
||
// ScreenshotHTMLToMessageChain 对HTML内容进行截图并返回包含图片的消息链
|
||
func ScreenshotHTMLToMessageChain(ctx context.Context, htmlContent string, opts *ScreenshotOptions) (protocol.MessageChain, error) {
|
||
base64Data, err := ScreenshotHTML(ctx, htmlContent, opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
chain := protocol.NewMessageChain()
|
||
chain = chain.AppendImageFromBase64(base64Data)
|
||
return chain, nil
|
||
}
|
||
|
||
// ScreenshotURLToMessageChain 对URL进行截图并返回包含图片的消息链
|
||
func ScreenshotURLToMessageChain(ctx context.Context, url string, opts *ScreenshotOptions) (protocol.MessageChain, error) {
|
||
base64Data, err := ScreenshotURL(ctx, url, opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
chain := protocol.NewMessageChain()
|
||
chain = chain.AppendImageFromBase64(base64Data)
|
||
return chain, nil
|
||
}
|
||
|
||
// ============================================================================
|
||
// HTML 模板功能
|
||
// ============================================================================
|
||
|
||
// RenderHTMLTemplate 渲染HTML模板并返回HTML字符串
|
||
func RenderHTMLTemplate(tmplContent string, data interface{}) (string, error) {
|
||
tmpl, err := template.New("html").Parse(tmplContent)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to parse template: %w", err)
|
||
}
|
||
|
||
var buf bytes.Buffer
|
||
if err := tmpl.Execute(&buf, data); err != nil {
|
||
return "", fmt.Errorf("failed to execute template: %w", err)
|
||
}
|
||
|
||
return buf.String(), nil
|
||
}
|
||
|
||
// RenderHTMLTemplateFromFile 从文件加载并渲染HTML模板
|
||
func RenderHTMLTemplateFromFile(tmplPath string, data interface{}) (string, error) {
|
||
content, err := os.ReadFile(tmplPath)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to read template file: %w", err)
|
||
}
|
||
|
||
tmpl, err := template.New(filepath.Base(tmplPath)).Parse(string(content))
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to parse template: %w", err)
|
||
}
|
||
|
||
var buf bytes.Buffer
|
||
if err := tmpl.Execute(&buf, data); err != nil {
|
||
return "", fmt.Errorf("failed to execute template: %w", err)
|
||
}
|
||
|
||
return buf.String(), nil
|
||
}
|
||
|
||
// ScreenshotTemplate 渲染HTML模板并截图，返回base64编码
|
||
func ScreenshotTemplate(ctx context.Context, tmplContent string, data interface{}, opts *ScreenshotOptions) (string, error) {
|
||
html, err := RenderHTMLTemplate(tmplContent, data)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
return ScreenshotHTML(ctx, html, opts)
|
||
}
|
||
|
||
// ScreenshotTemplateFromFile 从文件加载模板，渲染并截图，返回base64编码
|
||
func ScreenshotTemplateFromFile(ctx context.Context, tmplPath string, data interface{}, opts *ScreenshotOptions) (string, error) {
|
||
html, err := RenderHTMLTemplateFromFile(tmplPath, data)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
return ScreenshotHTML(ctx, html, opts)
|
||
}
|
||
|
||
// ScreenshotTemplateToMessageChain 渲染HTML模板并截图，返回包含图片的消息链
|
||
func ScreenshotTemplateToMessageChain(ctx context.Context, tmplContent string, data interface{}, opts *ScreenshotOptions) (protocol.MessageChain, error) {
|
||
base64Data, err := ScreenshotTemplate(ctx, tmplContent, data, opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
chain := protocol.NewMessageChain()
|
||
chain = chain.AppendImageFromBase64(base64Data)
|
||
return chain, nil
|
||
}
|
||
|
||
// ScreenshotTemplateFromFileToMessageChain 从文件加载模板，渲染并截图，返回包含图片的消息链
|
||
func ScreenshotTemplateFromFileToMessageChain(ctx context.Context, tmplPath string, data interface{}, opts *ScreenshotOptions) (protocol.MessageChain, error) {
|
||
base64Data, err := ScreenshotTemplateFromFile(ctx, tmplPath, data, opts)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
chain := protocol.NewMessageChain()
|
||
chain = chain.AppendImageFromBase64(base64Data)
|
||
return chain, nil
|
||
}
|