1package agents
2
3import (
4	"context"
5	"fmt"
6	"io/ioutil"
7	"os"
8	"os/exec"
9	"regexp"
10	"strconv"
11	"strings"
12	"time"
13
14	"github.com/michenriksen/aquatone/core"
15)
16
17type URLScreenshotter struct {
18	session         *core.Session
19	chromePath      string
20	tempUserDirPath string
21}
22
23func NewURLScreenshotter() *URLScreenshotter {
24	return &URLScreenshotter{}
25}
26
27func (a *URLScreenshotter) ID() string {
28	return "agent:url_screenshotter"
29}
30
31func (a *URLScreenshotter) Register(s *core.Session) error {
32	s.EventBus.SubscribeAsync(core.URLResponsive, a.OnURLResponsive, false)
33	s.EventBus.SubscribeAsync(core.SessionEnd, a.OnSessionEnd, false)
34	a.session = s
35	a.createTempUserDir()
36	a.locateChrome()
37
38	return nil
39}
40
41func (a *URLScreenshotter) OnURLResponsive(url string) {
42	a.session.Out.Debug("[%s] Received new responsive URL %s\n", a.ID(), url)
43	page := a.session.GetPage(url)
44	if page == nil {
45		a.session.Out.Error("Unable to find page for URL: %s\n", url)
46		return
47	}
48
49	a.session.WaitGroup.Add()
50	go func(page *core.Page) {
51		defer a.session.WaitGroup.Done()
52		a.screenshotPage(page)
53	}(page)
54}
55
56func (a *URLScreenshotter) OnSessionEnd() {
57	a.session.Out.Debug("[%s] Received SessionEnd event\n", a.ID())
58	os.RemoveAll(a.tempUserDirPath)
59	a.session.Out.Debug("[%s] Deleted temporary user directory at: %s\n", a.ID(), a.tempUserDirPath)
60}
61
62func (a *URLScreenshotter) createTempUserDir() {
63	dir, err := ioutil.TempDir("", "aquatone-chrome")
64	if err != nil {
65		a.session.Out.Fatal("Unable to create temporary user directory for Chrome/Chromium browser\n")
66		os.Exit(1)
67	}
68	a.session.Out.Debug("[%s] Created temporary user directory at: %s\n", a.ID(), dir)
69	a.tempUserDirPath = dir
70}
71
72func (a *URLScreenshotter) locateChrome() {
73	if *a.session.Options.ChromePath != "" {
74		a.chromePath = *a.session.Options.ChromePath
75		return
76	}
77
78	paths := []string{
79		"/usr/bin/google-chrome",
80		"/usr/bin/google-chrome-beta",
81		"/usr/bin/google-chrome-unstable",
82		"/usr/bin/chromium-browser",
83		"/usr/bin/chromium",
84		"/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
85		"/Applications/Google Chrome Canary.app/Contents/MacOS/Google Chrome Canary",
86		"/Applications/Chromium.app/Contents/MacOS/Chromium",
87		"C:/Program Files (x86)/Google/Chrome/Application/chrome.exe",
88	}
89
90	for _, path := range paths {
91		if _, err := os.Stat(path); os.IsNotExist(err) {
92			continue
93		}
94		a.chromePath = path
95	}
96
97	if a.chromePath == "" {
98		a.session.Out.Fatal("Unable to locate a valid installation of Chrome. Install Google Chrome or try specifying a valid location with the -chrome-path option.\n")
99		os.Exit(1)
100	}
101
102	if strings.Contains(strings.ToLower(a.chromePath), "chrome") {
103		a.session.Out.Warn("Using unreliable Google Chrome for screenshots. Install Chromium for better results.\n\n")
104	} else {
105		out, err := exec.Command(a.chromePath, "--version").Output()
106		if err != nil {
107			a.session.Out.Warn("An error occurred while trying to determine version of Chromium.\n\n")
108			return
109		}
110		version := string(out)
111		re := regexp.MustCompile(`(\d+)\.`)
112		match := re.FindStringSubmatch(version)
113		if len(match) <= 0 {
114			a.session.Out.Warn("Unable to determine version of Chromium. Screenshotting might be unreliable.\n\n")
115			return
116		}
117		majorVersion, _ := strconv.Atoi(match[1])
118		if majorVersion < 72 {
119			a.session.Out.Warn("An older version of Chromium is installed. Screenshotting of HTTPS URLs might be unreliable.\n\n")
120		}
121	}
122
123	a.session.Out.Debug("[%s] Located Chrome/Chromium binary at %s\n", a.ID(), a.chromePath)
124}
125
126func (a *URLScreenshotter) screenshotPage(page *core.Page) {
127	filePath := fmt.Sprintf("screenshots/%s.png", page.BaseFilename())
128	var chromeArguments = []string{
129		"--headless", "--disable-gpu", "--hide-scrollbars", "--mute-audio", "--disable-notifications",
130		"--no-first-run", "--disable-crash-reporter", "--ignore-certificate-errors", "--incognito",
131		"--disable-infobars", "--disable-sync", "--no-default-browser-check",
132		"--user-data-dir=" + a.tempUserDirPath,
133		"--user-agent=" + RandomUserAgent(),
134		"--window-size=" + *a.session.Options.Resolution,
135		"--screenshot=" + a.session.GetFilePath(filePath),
136	}
137
138	if os.Geteuid() == 0 {
139		chromeArguments = append(chromeArguments, "--no-sandbox")
140	}
141
142	if *a.session.Options.Proxy != "" {
143		chromeArguments = append(chromeArguments, "--proxy-server="+*a.session.Options.Proxy)
144	}
145
146	chromeArguments = append(chromeArguments, page.URL)
147
148	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(*a.session.Options.ScreenshotTimeout)*time.Millisecond)
149	defer cancel()
150
151	cmd := exec.CommandContext(ctx, a.chromePath, chromeArguments...)
152	if err := cmd.Start(); err != nil {
153		a.session.Out.Debug("[%s] Error: %v\n", a.ID(), err)
154		a.session.Stats.IncrementScreenshotFailed()
155		a.session.Out.Error("%s: screenshot failed: %s\n", page.URL, err)
156		a.killChromeProcessIfRunning(cmd)
157		return
158	}
159
160	if err := cmd.Wait(); err != nil {
161		a.session.Stats.IncrementScreenshotFailed()
162		a.session.Out.Debug("[%s] Error: %v\n", a.ID(), err)
163		if ctx.Err() == context.DeadlineExceeded {
164			a.session.Out.Error("%s: screenshot timed out\n", page.URL)
165			a.killChromeProcessIfRunning(cmd)
166			return
167		}
168
169		a.session.Out.Error("%s: screenshot failed: %s\n", page.URL, err)
170		a.killChromeProcessIfRunning(cmd)
171		return
172	}
173
174	a.session.Stats.IncrementScreenshotSuccessful()
175	a.session.Out.Info("%s: %s\n", page.URL, Green("screenshot successful"))
176	page.ScreenshotPath = filePath
177	page.HasScreenshot = true
178	a.killChromeProcessIfRunning(cmd)
179}
180
181func (a *URLScreenshotter) killChromeProcessIfRunning(cmd *exec.Cmd) {
182	if cmd.Process == nil {
183		return
184	}
185	cmd.Process.Release()
186	cmd.Process.Kill()
187}
188