143 lines
4.7 KiB
Go
143 lines
4.7 KiB
Go
package util
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/chromedp/cdproto/emulation"
|
|
"github.com/chromedp/cdproto/network"
|
|
"github.com/chromedp/cdproto/page"
|
|
"github.com/chromedp/chromedp"
|
|
)
|
|
|
|
// ScreenshotURL 截图函数：传入网址、输出路径、宽高、四个边距和等待的元素 ID
|
|
func ScreenshotURL(url, output string, width, height int, marginTop, marginBottom, marginRight, marginLeft int, waitClass string) error {
|
|
|
|
// 创建一个上下文，连接到 Docker 中运行的 headless-shell 实例
|
|
remoteAllocatorCtx, cancel := chromedp.NewRemoteAllocator(
|
|
context.Background(), "ws://100.124.180.117:9222/json/ws",
|
|
)
|
|
defer cancel()
|
|
|
|
ctx, cancel := chromedp.NewContext(remoteAllocatorCtx)
|
|
defer cancel()
|
|
|
|
// 设置超时时间，避免长时间无响应
|
|
ctx, cancel = context.WithTimeout(ctx, 60*time.Second)
|
|
defer cancel()
|
|
|
|
// 设置页面的宽高和缩放
|
|
if err := chromedp.Run(ctx, setViewport(width, height)); err != nil {
|
|
return fmt.Errorf("设置页面大小失败: %w", err)
|
|
}
|
|
|
|
// 启用网络请求拦截
|
|
if err := chromedp.Run(ctx, enableRequestInterception()); err != nil {
|
|
return fmt.Errorf("启用请求拦截失败: %w", err)
|
|
}
|
|
|
|
// 用于存储截图的变量
|
|
var screenshot []byte
|
|
|
|
queryAction := chromedp.WaitVisible(fmt.Sprintf(".%s", waitClass), chromedp.ByQuery)
|
|
if waitClass == "" {
|
|
queryAction = chromedp.WaitVisible(`body`, chromedp.ByQuery)
|
|
}
|
|
// 设置自定义请求头
|
|
headers := map[string]interface{}{
|
|
"accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
|
|
"accept-language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
|
|
"cache-control": "max-age=0",
|
|
"priority": "u=0, i",
|
|
"sec-ch-ua": `"Microsoft Edge";v="129", "Not=A?Brand";v="8", "Chromium";v="129"`,
|
|
"sec-ch-ua-mobile": "?0",
|
|
"sec-ch-ua-platform": `"Windows"`,
|
|
"sec-fetch-dest": "document",
|
|
"sec-fetch-mode": "navigate",
|
|
"sec-fetch-site": "same-origin",
|
|
"sec-fetch-user": "?1",
|
|
"upgrade-insecure-requests": "1",
|
|
"user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36 Edg/129.0.0.0",
|
|
}
|
|
|
|
// 启用网络功能并设置请求头
|
|
if err := chromedp.Run(ctx, network.Enable()); err != nil {
|
|
return fmt.Errorf("启用网络失败: %w", err)
|
|
}
|
|
if err := chromedp.Run(ctx, network.SetExtraHTTPHeaders(network.Headers(headers))); err != nil {
|
|
return fmt.Errorf("设置请求头失败: %w", err)
|
|
}
|
|
// 执行任务：打开网页并截图
|
|
err := chromedp.Run(ctx,
|
|
chromedp.Navigate(url), // 打开网页
|
|
ignoreErrors(queryAction), // 等待指定元素
|
|
chromedp.Sleep(10*time.Second),
|
|
chromedp.ActionFunc(func(ctx context.Context) error { // 自定义截图逻辑
|
|
_, _, _, _, cssLayoutViewport, _, _ := page.GetLayoutMetrics().Do(ctx)
|
|
if width == 0 || height == 0 {
|
|
width = int(cssLayoutViewport.ClientWidth)
|
|
height = int(cssLayoutViewport.ClientHeight)
|
|
}
|
|
// 计算调整后的截图区域
|
|
clip := &page.Viewport{
|
|
X: float64(marginLeft),
|
|
Y: float64(marginTop),
|
|
Width: float64(width - marginLeft - marginRight),
|
|
Height: float64(height - marginTop - marginBottom),
|
|
Scale: 1.0,
|
|
}
|
|
var err error
|
|
screenshot, err = page.CaptureScreenshot().WithClip(clip).Do(ctx)
|
|
return err
|
|
}),
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("截图失败: %w", err)
|
|
}
|
|
|
|
// 保存截图到本地
|
|
if err := os.WriteFile(output, screenshot, 0644); err != nil {
|
|
return fmt.Errorf("保存图片失败: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func setViewport(width, height int) chromedp.Tasks {
|
|
return chromedp.Tasks{
|
|
emulation.SetDeviceMetricsOverride(int64(width), int64(height), 1.0, false).
|
|
WithScreenOrientation(&emulation.ScreenOrientation{
|
|
Type: emulation.OrientationTypePortraitPrimary,
|
|
Angle: 0,
|
|
}),
|
|
}
|
|
}
|
|
|
|
// enableRequestInterception 启用网络请求拦截，忽略不必要的资源加载（如广告、图片等）
|
|
func enableRequestInterception() chromedp.Tasks {
|
|
return chromedp.Tasks{
|
|
chromedp.ActionFunc(func(ctx context.Context) error {
|
|
return network.Enable().Do(ctx)
|
|
}),
|
|
network.SetBlockedURLS([]string{
|
|
"pagead2.googlesyndication.com",
|
|
"optimizationguide-pa.googleapis.com",
|
|
"static.cloudflareinsights.com",
|
|
}),
|
|
}
|
|
}
|
|
|
|
// ignoreErrors 包裹 chromedp 任务，忽略执行过程中出现的错误
|
|
func ignoreErrors(task chromedp.Action) chromedp.ActionFunc {
|
|
return chromedp.ActionFunc(func(ctx context.Context) error {
|
|
err := task.Do(ctx)
|
|
if err != nil {
|
|
log.Printf("忽略错误: %v", err)
|
|
}
|
|
return nil
|
|
})
|
|
}
|