Repository: raviqqe/muffet
Pull request #220 (2022-05-17 12:38:46)

feature: added include pattern


Patch:
diff --git a/arguments.go b/arguments.go
index 4e10452..62e97de 100644
--- a/arguments.go
+++ b/arguments.go
@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"errors"
+	"fmt"
 	"regexp"
 	"strings"
 
@@ -14,6 +15,7 @@ type arguments struct {
 	MaxConnections        int      `short:"c" long:"max-connections" value-name:"<count>" default:"512" description:"Maximum number of HTTP connections"`
 	MaxConnectionsPerHost int      `long:"max-connections-per-host" value-name:"<count>" default:"512" description:"Maximum number of HTTP connections per host"`
 	RawExcludedPatterns   []string `short:"e" long:"exclude" value-name:"<pattern>..." description:"Exclude URLs matched with given regular expressions"`
+	RawIncludePatterns    []string `short:"i" long:"include" value-name:"<pattern>..." description:"Include URLs matched with given regular expressions"`
 	FollowRobotsTxt       bool     `long:"follow-robots-txt" description:"Follow robots.txt when scraping pages"`
 	FollowSitemapXML      bool     `long:"follow-sitemap-xml" description:"Scrape only pages listed in sitemap.xml"`
 	RawHeaders            []string `long:"header" value-name:"<header>..." description:"Custom headers"`
@@ -31,6 +33,7 @@ type arguments struct {
 	Version               bool     `long:"version" description:"Show version"`
 	URL                   string
 	ExcludedPatterns      []*regexp.Regexp
+	IncludePatterns       []*regexp.Regexp
 	Headers               map[string]string
 }
 
@@ -48,12 +51,19 @@ func getArguments(ss []string) (*arguments, error) {
 
 	args.URL = ss[0]
 
-	rs, err := compileRegexps(args.RawExcludedPatterns)
+	res, err := compileRegexps(args.RawExcludedPatterns)
 	if err != nil {
 		return nil, err
 	}
 
-	args.ExcludedPatterns = rs
+	args.ExcludedPatterns = res
+
+	ris, err := compileRegexps(args.RawIncludePatterns)
+	if err != nil {
+		return nil, errors.New(fmt.Sprintf("Parse include patterns: %s", err.Error()))
+	}
+
+	args.IncludePatterns = ris
 
 	hs, err := parseHeaders(args.RawHeaders)
 	if err != nil {
diff --git a/command.go b/command.go
index ae4768b..ac6f967 100644
--- a/command.go
+++ b/command.go
@@ -64,7 +64,7 @@ func (c *command) runWithError(ss []string) (bool, error) {
 		args.MaxRedirections,
 	)
 
-	pp := newPageParser(newLinkFinder(args.ExcludedPatterns))
+	pp := newPageParser(newLinkFinder(args.ExcludedPatterns, args.IncludePatterns))
 
 	f := newLinkFetcher(
 		client,
diff --git a/link_finder.go b/link_finder.go
index 4103f5a..c081717 100644
--- a/link_finder.go
+++ b/link_finder.go
@@ -31,10 +31,11 @@ var imageDescriptorPattern = regexp.MustCompile(" [^ ]*$")
 
 type linkFinder struct {
 	excludedPatterns []*regexp.Regexp
+	includedPatterns []*regexp.Regexp
 }
 
-func newLinkFinder(rs []*regexp.Regexp) linkFinder {
-	return linkFinder{rs}
+func newLinkFinder(rs []*regexp.Regexp, includePatterns []*regexp.Regexp) linkFinder {
+	return linkFinder{excludedPatterns: rs, includedPatterns: includePatterns}
 }
 
 func (f linkFinder) Find(n *html.Node, base *url.URL) map[string]error {
@@ -54,6 +55,11 @@ func (f linkFinder) Find(n *html.Node, base *url.URL) map[string]error {
 					continue
 				}
 
+				// only use include pattern when not empty
+				if len(f.includedPatterns) > 0 && !f.isLinkIncluded(s) {
+					continue
+				}
+
 				u, err := url.Parse(s)
 				if err != nil {
 					ls[s] = err
@@ -84,7 +90,15 @@ func (linkFinder) parseLinks(n *html.Node, a string) []string {
 }
 
 func (f linkFinder) isLinkExcluded(u string) bool {
-	for _, r := range f.excludedPatterns {
+	return f.isMatch(u, f.excludedPatterns)
+}
+
+func (f linkFinder) isLinkIncluded(u string) bool {
+	return f.isMatch(u, f.includedPatterns)
+}
+
+func (f linkFinder) isMatch(u string, patterns []*regexp.Regexp) bool {
+	for _, r := range patterns {
 		if r.MatchString(u) {
 			return true
 		}


Test patch:
diff --git a/.snapshots/TestHelp b/.snapshots/TestHelp
index 55e127f..40295f0 100644
--- a/.snapshots/TestHelp
+++ b/.snapshots/TestHelp
@@ -10,6 +10,8 @@ Application Options:
                                             per host (default: 512)
   -e, --exclude=<pattern>...                Exclude URLs matched with given
                                             regular expressions
+  -i, --include=<pattern>...                Include URLs matched with given
+                                            regular expressions
       --follow-robots-txt                   Follow robots.txt when scraping
                                             pages
       --follow-sitemap-xml                  Scrape only pages listed in
diff --git a/link_fetcher_test.go b/link_fetcher_test.go
index 3c2a822..8695f8e 100644
--- a/link_fetcher_test.go
+++ b/link_fetcher_test.go
@@ -15,7 +15,7 @@ func newTestLinkFetcher(c *fakeHttpClient) *linkFetcher {
 }
 
 func newTestLinkFetcherWithOptions(c *fakeHttpClient, o linkFetcherOptions) *linkFetcher {
-	return newLinkFetcher(c, newPageParser(newLinkFinder(nil)), o)
+	return newLinkFetcher(c, newPageParser(newLinkFinder(nil, nil)), o)
 }
 
 func TestNewFetcher(t *testing.T) {
diff --git a/link_finder_test.go b/link_finder_test.go
index 66159c8..3e723c4 100644
--- a/link_finder_test.go
+++ b/link_finder_test.go
@@ -38,7 +38,7 @@ func TestLinkFinderFindLinks(t *testing.T) {
 
 		s, e := 0, 0
 
-		for _, err := range newLinkFinder(nil).Find(n, b) {
+		for _, err := range newLinkFinder(nil, nil).Find(n, b) {
 			if err == nil {
 				s++
 			} else {
@@ -60,7 +60,7 @@ func TestLinkFinderFindLinkWithoutEncodedSpaces(t *testing.T) {
 	)
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	err, ok := ls["http://foo.com/a%20b"]
 	assert.True(t, ok)
@@ -76,7 +76,7 @@ func TestLinkFinderFindLinkWithoutSpacesNotEncoded(t *testing.T) {
 	)
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	err, ok := ls["http://foo.com/a%20b"]
 	assert.True(t, ok)
@@ -92,7 +92,7 @@ func TestLinkFinderFindLinkWithLeadingAndTrailingSpaces(t *testing.T) {
 	)
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	err, ok := ls["http://foo.com"]
 	assert.True(t, ok)
@@ -106,7 +106,7 @@ func TestLinkFinderFailWithInvalidURL(t *testing.T) {
 	n, err := html.Parse(strings.NewReader(htmlWithBody(`<a href=":" />`)))
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	assert.Equal(t, 1, len(ls))
 	assert.NotNil(t, ls[":"])
@@ -137,7 +137,36 @@ func TestLinkFinderIsLinkExcluded(t *testing.T) {
 		rs, err := compileRegexps(x.regexps)
 		assert.Nil(t, err)
 
-		assert.Equal(t, x.answer, newLinkFinder(rs).isLinkExcluded("http://foo.com"))
+		assert.Equal(t, x.answer, newLinkFinder(rs, nil).isLinkExcluded("http://foo.com"))
+	}
+}
+
+func TestLinkFinderIsLinkIncluded(t *testing.T) {
+	for _, x := range []struct {
+		regexps []string
+		answer  bool
+	}{
+		{
+			[]string{"foo\\.com"},
+			true,
+		},
+		{
+			[]string{"foo"},
+			true,
+		},
+		{
+			[]string{"bar", "foo"},
+			true,
+		},
+		{
+			[]string{"bar"},
+			false,
+		},
+	} {
+		rs, err := compileRegexps(x.regexps)
+		assert.Nil(t, err)
+
+		assert.Equal(t, x.answer, newLinkFinder(nil, rs).isLinkIncluded("http://foo.com"))
 	}
 }
 
@@ -150,7 +179,7 @@ func TestLinkFinderFindLinkInSrcSet(t *testing.T) {
 	)
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	err, ok := ls["http://foo.com/foo.png"]
 	assert.True(t, ok)
@@ -166,7 +195,7 @@ func TestLinkFinderFindMultipleLinksInSrcSet(t *testing.T) {
 	)
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	err, ok := ls["http://foo.com/foo.png"]
 	assert.True(t, ok)
@@ -186,7 +215,7 @@ func TestLinkFinderFindMultipleLinksInSrcSetWithDescriptors(t *testing.T) {
 	)
 	assert.Nil(t, err)
 
-	ls := newLinkFinder(nil).Find(n, b)
+	ls := newLinkFinder(nil, nil).Find(n, b)
 
 	err, ok := ls["http://foo.com/foo.png"]
 	assert.True(t, ok)
diff --git a/page_checker_test.go b/page_checker_test.go
index 2333d29..199c6fe 100644
--- a/page_checker_test.go
+++ b/page_checker_test.go
@@ -12,7 +12,7 @@ func newTestPageChecker(c *fakeHttpClient) *pageChecker {
 	return newPageChecker(
 		newLinkFetcher(
 			c,
-			newPageParser(newLinkFinder(nil)),
+			newPageParser(newLinkFinder(nil, nil)),
 			linkFetcherOptions{},
 		),
 		newLinkValidator("foo.com", nil, nil),
diff --git a/page_parser_test.go b/page_parser_test.go
index a475942..370c049 100644
--- a/page_parser_test.go
+++ b/page_parser_test.go
@@ -7,19 +7,19 @@ import (
 )
 
 func TestPageParserParsePage(t *testing.T) {
-	_, err := newPageParser(newLinkFinder(nil)).Parse("http://foo.com", nil)
+	_, err := newPageParser(newLinkFinder(nil, nil)).Parse("http://foo.com", nil)
 	assert.Nil(t, err)
 }
 
 func TestPageParserFailWithInvalidURL(t *testing.T) {
-	_, err := newPageParser(newLinkFinder(nil)).Parse(":", nil)
+	_, err := newPageParser(newLinkFinder(nil, nil)).Parse(":", nil)
 	assert.NotNil(t, err)
 }
 
 func TestPageParserSetCorrectURL(t *testing.T) {
 	s := "http://foo.com"
 
-	p, err := newPageParser(newLinkFinder(nil)).Parse(s, nil)
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(s, nil)
 	assert.Nil(t, err)
 	assert.Equal(t, s, p.URL().String())
 }
@@ -27,7 +27,7 @@ func TestPageParserSetCorrectURL(t *testing.T) {
 func TestPageParserIgnorePageURLFragment(t *testing.T) {
 	s := "http://foo.com#id"
 
-	p, err := newPageParser(newLinkFinder(nil)).Parse(s, nil)
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(s, nil)
 	assert.Nil(t, err)
 	assert.Equal(t, "http://foo.com", p.URL().String())
 }
@@ -35,13 +35,13 @@ func TestPageParserIgnorePageURLFragment(t *testing.T) {
 func TestPageParserKeepQuery(t *testing.T) {
 	s := "http://foo.com?bar=baz"
 
-	p, err := newPageParser(newLinkFinder(nil)).Parse(s, nil)
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(s, nil)
 	assert.Nil(t, err)
 	assert.Equal(t, s, p.URL().String())
 }
 
 func TestPageParserResolveLinksWithBaseTag(t *testing.T) {
-	p, err := newPageParser(newLinkFinder(nil)).Parse(
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 		"http://foo.com",
 		[]byte(`
 			<html>
@@ -59,7 +59,7 @@ func TestPageParserResolveLinksWithBaseTag(t *testing.T) {
 }
 
 func TestPageParserResolveLinksWithBlankBaseTag(t *testing.T) {
-	p, err := newPageParser(newLinkFinder(nil)).Parse(
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 		"http://foo.com",
 		[]byte(`
 			<html>
@@ -77,7 +77,7 @@ func TestPageParserResolveLinksWithBlankBaseTag(t *testing.T) {
 }
 
 func TestPageParserFailToParseWithInvalidBaseTag(t *testing.T) {
-	_, err := newPageParser(newLinkFinder(nil)).Parse(
+	_, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 		"http://foo.com",
 		[]byte(`
 			<html>
@@ -93,7 +93,7 @@ func TestPageParserFailToParseWithInvalidBaseTag(t *testing.T) {
 }
 
 func TestPageParserParseID(t *testing.T) {
-	p, err := newPageParser(newLinkFinder(nil)).Parse(
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 		"http://foo.com",
 		[]byte(`<p id="foo" />`),
 	)
@@ -102,7 +102,7 @@ func TestPageParserParseID(t *testing.T) {
 }
 
 func TestPageParserParseName(t *testing.T) {
-	p, err := newPageParser(newLinkFinder(nil)).Parse(
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 		"http://foo.com",
 		[]byte(`<p name="foo" />`),
 	)
@@ -111,7 +111,7 @@ func TestPageParserParseName(t *testing.T) {
 }
 
 func TestPageParserParseIDAndName(t *testing.T) {
-	p, err := newPageParser(newLinkFinder(nil)).Parse(
+	p, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 		"http://foo.com",
 		[]byte(`<p id="foo" name="bar" />`),
 	)
@@ -130,7 +130,7 @@ func TestPageParserParseLinks(t *testing.T) {
 			"http://foo.com/foo.img",
 		},
 	} {
-		p, err := newPageParser(newLinkFinder(nil)).Parse(
+		p, err := newPageParser(newLinkFinder(nil, nil)).Parse(
 			"http://foo.com",
 			[]byte(ss[0]),
 		)