Skip to content

Commit

Permalink
Refactor: accept single url
Browse files Browse the repository at this point in the history
  • Loading branch information
@web-flow
web-flow committed Jul 3, 2021
1 parent c4aef71 commit 45d08b9
Show file tree
Hide file tree
Showing 3 changed files with 192 additions and 192 deletions.
95 changes: 49 additions & 46 deletions cmd/screenshot/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@ import (
"fmt"
"io/ioutil"
"net/http"
"net/url"
"os"
"sync"
"time"

"github.com/wabarc/helper"
Expand Down Expand Up @@ -42,64 +44,65 @@ func main() {
os.Exit(1)
}

ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
defer cancel()
var write = func(uri string, data []byte) {
if data == nil {
return
}

var str string
for _, arg := range args {
str += fmt.Sprintf(" %s", arg)
filename := helper.FileName(uri, http.DetectContentType(data))
if err := ioutil.WriteFile(filename, data, 0o644); err != nil {
fmt.Println(uri, "=>", err)
return
}
fmt.Println(uri, "=>", filename)
}

urls := helper.MatchURL(str)
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
defer cancel()

var err error
var shots []screenshot.Screenshots
var shot screenshot.Screenshots
var opts = []screenshot.ScreenshotOption{
screenshot.ScaleFactor(1),
screenshot.PrintPDF(pdf), // print pdf
screenshot.RawHTML(raw), // export html
screenshot.Quality(100), // image quality
}
if remoteAddr != "" {
remote, er := screenshot.NewChromeRemoteScreenshoter(remoteAddr)
if er != nil {
fmt.Println(er)
return
}
shots, err = remote.Screenshot(ctx, urls, opts...)
} else {
shots, err = screenshot.Screenshot(ctx, urls, opts...)
}
if err != nil {
if err == context.DeadlineExceeded {
fmt.Println(err.Error())
return
}
fmt.Println(err.Error())
return
}

for _, shot := range shots {
if shot.URL == "" || shot.Image == nil {
continue
}
if img {
if err := ioutil.WriteFile(helper.FileName(shot.URL, http.DetectContentType(shot.Image)), shot.Image, 0o644); err != nil {
fmt.Println(err)
continue
var wg sync.WaitGroup
for _, arg := range args {
wg.Add(1)
go func(link string) {
defer wg.Done()
input, err := url.Parse(link)
if err != nil {
fmt.Println(link, "=>", fmt.Sprintf("%v", err))
return
}
}
if pdf {
if err := ioutil.WriteFile(helper.FileName(shot.URL, http.DetectContentType(shot.PDF)), shot.PDF, 0o644); err != nil {
fmt.Println(err)
continue
if remoteAddr != "" {
remote, er := screenshot.NewChromeRemoteScreenshoter(remoteAddr)
if er != nil {
fmt.Println(er)
return
}
shot, err = remote.Screenshot(ctx, input, opts...)
} else {
shot, err = screenshot.Screenshot(ctx, input, opts...)
}
}
if raw {
if err := ioutil.WriteFile(helper.FileName(shot.URL, http.DetectContentType(shot.HTML)), shot.HTML, 0o644); err != nil {
fmt.Println(err)
continue
if err != nil {
if err == context.DeadlineExceeded {
fmt.Println(err.Error())
return
}
fmt.Println(err.Error())
return
}
}

if shot.URL == "" || shot.Image == nil {
return
}
write(shot.URL, shot.Image)
write(shot.URL, shot.PDF)
write(shot.URL, shot.HTML)
}(arg)
}
wg.Wait()
}
153 changes: 74 additions & 79 deletions screenshot.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@ import (
"log"
"math"
"net/http"
"net/url"
"os"
"strings"
"sync"
"time"

"github.com/chromedp/cdproto/emulation"
Expand All @@ -36,7 +36,7 @@ type Screenshots struct {

// Screenshoter is a webpage screenshot interface.
type Screenshoter interface {
Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error)
Screenshot(ctx context.Context, input *url.URL, options ...ScreenshotOption) (Screenshots, error)
}

type chromeRemoteScreenshoter struct {
Expand Down Expand Up @@ -69,22 +69,22 @@ func NewChromeRemoteScreenshoter(addr string) (Screenshoter, error) {
}, nil
}

func (s *chromeRemoteScreenshoter) Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error) {
func (s *chromeRemoteScreenshoter) Screenshot(ctx context.Context, input *url.URL, options ...ScreenshotOption) (shot Screenshots, err error) {
if s.url == "" {
return nil, fmt.Errorf("Can't connect to headless browser")
return shot, fmt.Errorf("Can't connect to headless browser")
}

ctx, cancel := chromedp.NewRemoteAllocator(ctx, s.url)
defer cancel()

return screenshotStart(ctx, urls, options...)
return screenshotStart(ctx, input, options...)
}

func Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error) {
func Screenshot(ctx context.Context, input *url.URL, options ...ScreenshotOption) (Screenshots, error) {
// https://github.com/chromedp/chromedp/blob/b56cd66f9cebd6a1fa1283847bbf507409d48225/allocate.go#L53
var allocOpts = append(
chromedp.DefaultExecAllocatorOptions[:],
chromedp.CombinedOutput(log.Writer()),
// chromedp.CombinedOutput(log.Writer()),
chromedp.NoDefaultBrowserCheck,
chromedp.Flag("ignore-certificate-errors", true),
chromedp.Flag("allow-running-insecure-content", true),
Expand All @@ -103,10 +103,10 @@ func Screenshot(ctx context.Context, urls []string, options ...ScreenshotOption)
ctx, cancel := chromedp.NewExecAllocator(ctx, allocOpts...)
defer cancel()

return screenshotStart(ctx, urls, options...)
return screenshotStart(ctx, input, options...)
}

func screenshotStart(ctx context.Context, urls []string, options ...ScreenshotOption) ([]Screenshots, error) {
func screenshotStart(ctx context.Context, input *url.URL, options ...ScreenshotOption) (shot Screenshots, err error) {
var browserOpts []chromedp.ContextOption
if debug := os.Getenv("CHROMEDP_DEBUG"); debug != "" && debug != "false" {
browserOpts = append(browserOpts, chromedp.WithDebugf(log.Printf))
Expand All @@ -129,77 +129,72 @@ func screenshotStart(ctx context.Context, urls []string, options ...ScreenshotOp
// return nil, err
// }

var wg sync.WaitGroup
screenshots := make([]Screenshots, 0, len(urls))
for _, url := range urls {
wg.Add(1)
url := convertURI(url)
go func(url string) {
var buf []byte
var pdf []byte
var raw string
var title string
// var ok bool

chromedp.ListenTarget(ctx, func(ev interface{}) {
switch ev := ev.(type) {
case *page.EventJavascriptDialogOpening:
go func() {
if err := chromedp.Run(ctx,
page.HandleJavaScriptDialog(true),
); err != nil {
log.Print(err)
}
}()
// case *page.EventDocumentOpened:
// return
// case *network.EventRequestWillBeSent:
// return
// case *network.EventResponseReceived:
// return
case *network.EventLoadingFinished:
logger.Debug("[screenshot] EventLoadingFinished: %v", ev.RequestID)
return
url := convertURI(input.String())
var buf []byte
var pdf []byte
var raw string
var title string

chromedp.ListenTarget(ctx, func(ev interface{}) {
switch ev := ev.(type) {
case *page.EventJavascriptDialogOpening:
go func() {
if err := chromedp.Run(ctx,
page.HandleJavaScriptDialog(true),
); err != nil {
log.Print(err)
}
})

ctx, _ = chromedp.NewContext(ctx)
captureAction := screenshotAction(&buf, opts)
exportHTML := exportHTML(&raw, opts)
saveAsPDF := printPDF(&pdf, opts)
if err := chromedp.Run(ctx, chromedp.Tasks{
page.Enable(),
network.Enable(),
// enableLifeCycleEvents(),
page.SetDownloadBehavior(page.SetDownloadBehaviorBehaviorDeny),
navigateAndWaitFor(url, "networkIdle"),
// chromedp.Navigate(url),
chromedp.WaitReady("body"),
chromedp.Title(&title),
evaluate(nil),
captureAction,
exportHTML,
saveAsPDF,
chromedp.ResetViewport(),
chromedp.Sleep(time.Second),
closePageAction(),
}); err != nil {
log.Print(err)
buf = nil
}
screenshots = append(screenshots, Screenshots{
URL: revertURI(url),
PDF: pdf,
HTML: []byte(raw),
Image: buf,
Title: title,
})
wg.Done()
}(url)
}
wg.Wait()

return screenshots, nil
}()
// case *page.EventDocumentOpened:
// return
// case *network.EventRequestWillBeSent:
// return
// case *network.EventResponseReceived:
// return
case *network.EventLoadingFinished:
logger.Debug("[screenshot] EventLoadingFinished: %v", ev.RequestID)
return
}
})

ctx, _ = chromedp.NewContext(ctx)
captureAction := screenshotAction(&buf, opts)
exportHTML := exportHTML(&raw, opts)
saveAsPDF := printPDF(&pdf, opts)
if err := chromedp.Run(ctx, chromedp.Tasks{
page.Enable(),
network.Enable(),
// enableLifeCycleEvents(),
page.SetDownloadBehavior(page.SetDownloadBehaviorBehaviorDeny),
navigateAndWaitFor(url, "networkIdle"),
// chromedp.Navigate(url),
chromedp.WaitReady("body"),
chromedp.Title(&title),
evaluate(nil),
captureAction,
exportHTML,
saveAsPDF,
chromedp.ResetViewport(),
chromedp.Sleep(time.Second),
closePageAction(),
}); err != nil {
log.Print(err)
buf = nil
}

var html []byte
if raw != "" {
html = []byte(raw)
}
shot = Screenshots{
URL: revertURI(url),
PDF: pdf,
HTML: html,
Image: buf,
Title: title,
}

return shot, nil
}

func evaluate(res interface{}) chromedp.EvaluateAction {
Expand Down
Loading

0 comments on commit 45d08b9

Please sign in to comment.