refactor: use mapset instead of home grown set

lc · Jul 13, 2023 · 387822a · 387822a
1 parent 78e3d95
commit 387822a
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 48 deletions.
diff --git a/go.mod b/go.mod
@@ -1,9 +1,10 @@
 module github.com/lc/gau/v2
 
-go 1.17
+go 1.20
 
 require (
 	github.com/bobesa/go-domain-util v0.0.0-20190911083921-4033b5f7dd89
+	github.com/deckarep/golang-set/v2 v2.3.0
 	github.com/json-iterator/go v1.1.12
 	github.com/lynxsecurity/pflag v1.1.3
 	github.com/lynxsecurity/viper v1.10.0

diff --git a/go.sum b/go.sum
@@ -5,6 +5,8 @@ github.com/bobesa/go-domain-util v0.0.0-20190911083921-4033b5f7dd89/go.mod h1:/0
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/deckarep/golang-set/v2 v2.3.0 h1:qs18EKUfHm2X9fA50Mr/M5hccg2tNnVqsiBImnyDs0g=
+github.com/deckarep/golang-set/v2 v2.3.0/go.mod h1:VAky9rY/yGXJOLEDv3OMci+7wtDpOF4IN+y82NBOac4=
 github.com/fsnotify/fsnotify v1.5.1 h1:mZcQUHVQUQWoPXXtuf9yuEXKudkV2sx1E06UadKWpgI=
 github.com/fsnotify/fsnotify v1.5.1/go.mod h1:T3375wBYaZdLLcVNkcVbzGHY7f1l/uK5T5Ai1i3InKU=
 github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=

diff --git a/pkg/output/output.go b/pkg/output/output.go
@@ -1,6 +1,7 @@
 package output
 
 import (
+	mapset "github.com/deckarep/golang-set/v2"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/valyala/bytebufferpool"
 	"io"
@@ -13,40 +14,24 @@ type JSONResult struct {
 	Url string `json:"url"`
 }
 
-func WriteURLs(writer io.Writer, results <-chan string, blacklistMap map[string]struct{}, RemoveParameters bool) error {
-	lastURL := make(map[string]struct{})
+func WriteURLs(writer io.Writer, results <-chan string, blacklistMap mapset.Set[string], RemoveParameters bool) error {
+	lastURL := mapset.NewThreadUnsafeSet[string]()
 	for result := range results {
 		buf := bytebufferpool.Get()
-		if len(blacklistMap) != 0 {
-			u, err := url.Parse(result)
-			if err != nil {
-				continue
-			}
-			base := strings.Split(path.Base(u.Path), ".")
-			ext := base[len(base)-1]
-			if ext != "" {
-				_, ok := blacklistMap[strings.ToLower(ext)]
-				if ok {
-					continue
-				}
-			}
+		u, err := url.Parse(result)
+		if err != nil {
+			continue
 		}
-		if RemoveParameters {
-			u, err := url.Parse(result)
-			if err != nil {
-				continue
-			}
-			if _, ok := lastURL[u.Host+u.Path]; ok {
-				continue
-			} else {
-				lastURL[u.Host+u.Path] = struct{}{} ;
-			}
-
+		if blacklistMap.Contains(strings.ToLower(path.Ext(u.Path))) {
+			continue
+		}
+		if RemoveParameters && !lastURL.Add(u.Host+u.Path) {
+			continue
 		}
 
 		buf.B = append(buf.B, []byte(result)...)
 		buf.B = append(buf.B, "\n"...)
-		_, err := writer.Write(buf.B)
+		_, err = writer.Write(buf.B)
 		if err != nil {
 			return err
 		}
@@ -55,23 +40,16 @@ func WriteURLs(writer io.Writer, results <-chan string, blacklistMap map[string]
 	return nil
 }
 
-func WriteURLsJSON(writer io.Writer, results <-chan string, blacklistMap map[string]struct{}, RemoveParameters bool) {
+func WriteURLsJSON(writer io.Writer, results <-chan string, blacklistMap mapset.Set[string], RemoveParameters bool) {
 	var jr JSONResult
 	enc := jsoniter.NewEncoder(writer)
 	for result := range results {
-		if len(blacklistMap) != 0 {
-			u, err := url.Parse(result)
-			if err != nil {
-				continue
-			}
-			base := strings.Split(path.Base(u.Path), ".")
-			ext := base[len(base)-1]
-			if ext != "" {
-				_, ok := blacklistMap[strings.ToLower(ext)]
-				if ok {
-					continue
-				}
-			}
+		u, err := url.Parse(result)
+		if err != nil {
+			continue
+		}
+		if blacklistMap.Contains(strings.ToLower(path.Ext(u.Path))) {
+			continue
 		}
 		jr.Url = result
 		if err := enc.Encode(jr); err != nil {

diff --git a/pkg/providers/providers.go b/pkg/providers/providers.go
@@ -2,6 +2,7 @@ package providers
 
 import (
 	"context"
+	mapset "github.com/deckarep/golang-set/v2"
 	"github.com/valyala/fasthttp"
 )
 
@@ -27,7 +28,7 @@ type Config struct {
 	RemoveParameters  bool
 	Client            *fasthttp.Client
 	Providers         []string
-	Blacklist         map[string]struct{}
+	Blacklist         mapset.Set[string]
 	Output            string
 	JSON              bool
 	URLScan           URLScan

diff --git a/runner/flags/flags.go b/runner/flags/flags.go
@@ -4,6 +4,7 @@ import (
 	"crypto/tls"
 	"flag"
 	"fmt"
+	mapset "github.com/deckarep/golang-set/v2"
 	"github.com/lc/gau/v2/pkg/providers"
 	"github.com/lynxsecurity/pflag"
 	"github.com/lynxsecurity/viper"
@@ -80,11 +81,8 @@ func (c *Config) ProviderConfig() (*providers.Config, error) {
 		OTX: c.OTX,
 	}
 
-	pc.Blacklist = make(map[string]struct{})
-	for _, b := range c.Blacklist {
-		pc.Blacklist[b] = struct{}{}
-	}
-
+	pc.Blacklist = mapset.NewThreadUnsafeSet(c.Blacklist...)
+	pc.Blacklist.Add("")
 	return pc, nil
 }