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://127.0.0.1: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 }) }