Add custom scraper

.gitignore

@@ -0,0 +1 @@
+.idea

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/kennygrant/sanitize v1.2.4 // indirect
 	github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca // indirect
 	github.com/temoto/robotstxt v0.0.0-20180810133444-97ee4a9ee6ea // indirect
-	golang.org/x/net v0.0.0-20190522155817-f3200d17e092 // indirect
+	golang.org/x/net v0.0.0-20190522155817-f3200d17e092
 	golang.org/x/sync v0.0.0-20190423024810-112230192c58 // indirect
 	google.golang.org/appengine v1.6.0 // indirect
 )

main.go

@@ -4,24 +4,29 @@ import (
 	"flag"
 	"log"
 	"net/url"
+	"os"
 
 	"github.com/michaljemala/image-scraper/pkg/downloader"
 	"github.com/michaljemala/image-scraper/pkg/scraper"
 )
 
 var (
-	flagURL = flag.String("url", "https://www.google.com", "A website URL to be scraped")
-	flagDir = flag.String("dir", ".", "Destination directory where website assets will be scraped")
+	flagURL = flag.String("url", "https://exponea.com", "A website URL to be scraped")
+	flagDir = flag.String("dir", "./data", "Destination directory where website assets will be scraped")
 )
 
 func main() {
 	flag.Parse()
 
-	URL, err := url.Parse(*flagURL)
+	u, err := url.Parse(*flagURL)
 	if err != nil {
 		log.Fatal("invalid url")
 	}
 
+	if err := os.MkdirAll(*flagDir, os.ModePerm); err != nil {
+		log.Fatal("unable to create dir")
+	}
+
 	d, err := downloader.New(
 		downloader.WithDestDir(*flagDir),
 		downloader.WithConcurrency(10),
@@ -30,14 +35,21 @@ func main() {
 		log.Fatal(err)
 	}
 
-	s := scraper.NewScraper()
-	s.RegisterCallback("img[src]", func(e *scraper.HTMLElement) {
-		if url := e.AttrValue("src", true); url != "" {
-			d.Queue(url)
+	s := scraper.NewCustomScraper(scraper.Config{
+		Concurrency: 10,
+	})
+	s.OnHTMLElement("img", func(e *scraper.HTMLElement) {
+		link, ok := e.Attr("src")
+		if !ok {
+			return
+		}
+		link = e.ResolveURL(link)
+		if link != "" {
+			log.Printf("enqueueing: %s", link)
+			d.Queue(link)
 		}
 	})
-	s.Scrape(URL)
+	s.Scrape(u)
 
 	d.Stop()
 }
-

pkg/downloader/downloader.go

@@ -68,7 +68,7 @@ func (d *Downloader) DestDir() string {
 func (d *Downloader) Queue(url string) {
 	d.pool.Submit(worker.Task{
 		Info: url,
-		Fn: func () error {
+		Fn: func() error {
 			return d.download(url)
 		},
 	})
@@ -96,10 +96,16 @@ func (d *Downloader) download(url string) error {
 
 	resp, err := http.Get(url)
 	if err != nil {
-		return err
+		return fmt.Errorf("unable to get: %s: %v", url, err)
 	}
 	defer resp.Body.Close()
 
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("invalid status code: %s: %s", url, resp.Status)
+	}
+
+	log.Printf("downloaded: %s", url)
+
 	_, err = io.Copy(out, resp.Body)
 	return err
-}
+}

pkg/scraper/colly.go

@@ -0,0 +1,72 @@
+package scraper
+
+import (
+	"net/url"
+
+	"github.com/gocolly/colly"
+)
+
+type collyScraper struct {
+	cfg       Config
+	callbacks map[string]func(*HTMLElement)
+}
+
+func NewCollyScraper(cfg Config) *collyScraper {
+	if cfg.Concurrency <= 0 {
+		cfg.Concurrency = 10
+	}
+
+	return &collyScraper{
+		cfg:       cfg,
+		callbacks: make(map[string]func(*HTMLElement)),
+	}
+}
+
+func (s *collyScraper) OnHTMLElement(name string, callback func(*HTMLElement)) {
+	s.callbacks[name] = callback
+}
+
+func (s *collyScraper) Scrape(URL *url.URL) {
+	c := colly.NewCollector(
+		colly.AllowedDomains(URL.Hostname()),
+		colly.Async(true),
+	)
+
+	c.Limit(&colly.LimitRule{
+		DomainGlob:  "*",
+		Parallelism: s.cfg.Concurrency,
+	})
+
+	for name, fn := range s.callbacks {
+		c.OnHTML(name, func(e *colly.HTMLElement) {
+			if e == nil {
+				return
+			}
+
+			// TODO How we get *html.Node from *colly.HTMLElement
+			_ = fn
+			//n, err := html.Parse(strings.NewReader(e.Text))
+			//if err != nil {
+			//	log.Printf("unable to parse node: %s: %v", e.Name, err)
+			//	return
+			//}
+			//
+			//fn(&HTMLElement{
+			//	n:              n,
+			//	resolveURLFunc: e.Request.URL.Parse,
+			//})
+		})
+	}
+
+	c.OnHTML("a[href]", func(e *colly.HTMLElement) {
+		url := e.Request.AbsoluteURL(e.Attr("href"))
+		if url == "" {
+			return
+		}
+
+		c.Visit(url)
+	})
+
+	c.Visit(URL.String())
+	c.Wait()
+}

pkg/scraper/custom.go

@@ -0,0 +1,168 @@
+package scraper
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"net/url"
+	"sync"
+	"time"
+
+	"golang.org/x/net/html"
+)
+
+type customScraper struct {
+	config    Config
+	callbacks map[string]func(*HTMLElement)
+}
+
+func NewCustomScraper(cfg Config) *customScraper {
+	if cfg.Concurrency <= 0 {
+		cfg.Concurrency = 10
+	}
+	if cfg.Client == nil {
+		cfg.Client = http.DefaultClient
+	}
+	return &customScraper{
+		config:    cfg,
+		callbacks: make(map[string]func(*HTMLElement)),
+	}
+}
+
+func (s *customScraper) OnHTMLElement(name string, fn func(*HTMLElement)) {
+	s.callbacks[name] = fn
+}
+
+func (s *customScraper) Scrape(u *url.URL) {
+	seed := make(chan []string)
+	unseen := make(chan string)
+
+	allowedDomain := func(found *url.URL) bool {
+		return found.Hostname() == u.Hostname()
+	}
+
+	go func() {
+		seen := make(map[string]bool)
+		for list := range seed {
+			for _, link := range list {
+				if !seen[link] {
+					seen[link] = true
+					unseen <- link
+				}
+			}
+		}
+	}()
+
+	wg := &sync.WaitGroup{}
+	for i := 0; i < s.config.Concurrency; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			for {
+				select {
+				case <-time.After(1 * time.Second):
+					return // We have nothing to read for some time, we end !?!
+				case link := <-unseen:
+					doc, err := s.getDocument(link)
+					if err != nil {
+						s.handleError(link, err)
+						continue
+					}
+
+					resolveURL := func(ref string) (*url.URL, error) {
+						return doc.req.URL.Parse(ref)
+					}
+
+					var foundLinks []string
+					s.walkDocument(doc, func(n *html.Node) {
+						if n == nil || n.Type != html.ElementNode {
+							return
+						}
+
+						// Collect links
+						if n.Data == "a" {
+							for _, a := range n.Attr {
+								if a.Key != "href" {
+									continue
+								}
+								link, err := resolveURL(a.Val)
+								if err != nil {
+									continue
+								}
+								if allowedDomain(link) {
+									foundLinks = append(foundLinks, link.String())
+								}
+							}
+						}
+
+						// Call registered callbacks
+						if fn, ok := s.callbacks[n.Data]; ok {
+							fn(&HTMLElement{
+								n:              n,
+								resolveURLFunc: resolveURL,
+							})
+						}
+					})
+
+					if len(foundLinks) > 0 {
+						// Send the links asynchronously as we use unbuffered channel
+						go func() { seed <- foundLinks }()
+					}
+				}
+			}
+		}()
+	}
+
+	go func() { seed <- []string{u.String()} }()
+
+	wg.Wait()
+	close(seed)
+}
+
+type document struct {
+	root *html.Node
+	req  *http.Request
+}
+
+func (s *customScraper) getDocument(link string) (*document, error) {
+	resp, err := s.config.Client.Get(link)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("unable to get document: %s: %s", link, resp.Status)
+	}
+
+	doc, err := html.Parse(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse document: %s: %v", link, err)
+	}
+
+	return &document{
+		root: doc,
+		req:  resp.Request,
+	}, nil
+}
+
+func (s *customScraper) handleError(link string, err error) {
+	log.Printf("unable to load document: %s: %v", link, err)
+}
+
+func (s *customScraper) walkDocument(doc *document, fn func(n *html.Node)) {
+	if doc == nil || fn == nil {
+		return
+	}
+
+	walk(doc.root, fn)
+}
+
+func walk(n *html.Node, fn func(*html.Node)) {
+	fn(n)
+
+	for c := n.FirstChild; c != nil; c = c.NextSibling {
+		walk(c, fn)
+	}
+}