Merge pull request #15 from projectdiscovery/14-bugfix-adding-options-support

Adding support for auth options + stricter tests

Author: Ice3man543

cdncheck.go

@@ -12,19 +12,26 @@ import (
 // Client checks for CDN based IPs which should be excluded
 // during scans since they belong to third party firewalls.
 type Client struct {
-	Data   map[string]struct{}
-	ranger cidranger.Ranger
+	Options Options
+	Data    map[string]struct{}
+	ranger  cidranger.Ranger
 }
 
 var defaultScrapers = map[string]scraperFunc{
-	"akamai":     scrapeAkamai,
+	// "akamai":     scrapeAkamai,
 	"azure":      scrapeAzure,
 	"cloudflare": scrapeCloudflare,
 	"cloudfront": scrapeCloudFront,
 	"fastly":     scrapeFastly,
 	"incapsula":  scrapeIncapsula,
-	"sucuri":     scrapeSucuri,
-	"leaseweb":   scrapeLeaseweb,
+	// "sucuri":     scrapeSucuri,
+	// "leaseweb":   scrapeLeaseweb,
+}
+
+var defaultScrapersWithOptions = map[string]scraperWithOptionsFunc{
+	"akamai":   scrapeAkamai,
+	"sucuri":   scrapeSucuri,
+	"leaseweb": scrapeLeaseweb,
 }
 
 var cachedScrapers = map[string]scraperFunc{
@@ -33,15 +40,20 @@ var cachedScrapers = map[string]scraperFunc{
 
 // New creates a new firewall IP checking client.
 func New() (*Client, error) {
-	return new(false)
+	return new(&Options{})
 }
 
 // NewWithCache creates a new firewall IP with cached data from project discovery (faster)
 func NewWithCache() (*Client, error) {
-	return new(true)
+	return new(&Options{Cache: true})
+}
+
+// NewWithOptions creates a new instance with options
+func NewWithOptions(Options *Options) (*Client, error) {
+	return new(Options)
 }
 
-func new(cache bool) (*Client, error) {
+func new(options *Options) (*Client, error) {
 	httpClient := &http.Client{
 		Transport: &http.Transport{
 			MaxIdleConns:        100,
@@ -55,21 +67,31 @@ func new(cache bool) (*Client, error) {
 	}
 	client := &Client{}
 
-	var scrapers map[string]scraperFunc
-	if cache {
-		scrapers = cachedScrapers
+	if options.Cache {
+		for _, scraper := range cachedScrapers {
+			cidrs, err := scraper(httpClient)
+			if err != nil {
+				return nil, err
+			}
+			client.parseCidrs(cidrs)
+		}
 	} else {
-		scrapers = defaultScrapers
+		for _, scraper := range defaultScrapers {
+			cidrs, err := scraper(httpClient)
+			if err != nil {
+				return nil, err
+			}
+			client.parseCidrs(cidrs)
+		}
 	}
 
-	client.Data = make(map[string]struct{})
-	for _, scraper := range scrapers {
-		cidrs, err := scraper(httpClient)
-		if err != nil {
-			return nil, err
-		}
-		for _, cidr := range cidrs {
-			client.Data[cidr] = struct{}{}
+	if options.HasAuthInfo() {
+		for _, scraper := range defaultScrapersWithOptions {
+			cidrs, err := scraper(httpClient, options)
+			if err != nil {
+				return nil, err
+			}
+			client.parseCidrs(cidrs)
 		}
 	}
 
@@ -86,6 +108,16 @@ func new(cache bool) (*Client, error) {
 	return client, nil
 }
 
+// parseCidrs inserts the scraped cidrs to the internal structure
+func (c *Client) parseCidrs(cidrs []string) {
+	if c.Data == nil {
+		c.Data = make(map[string]struct{})
+	}
+	for _, cidr := range cidrs {
+		c.Data[cidr] = struct{}{}
+	}
+}
+
 // Check checks if an IP is contained in the blacklist
 func (c *Client) Check(ip net.IP) (bool, error) {
 	return c.ranger.Contains(ip)

options.go

@@ -0,0 +1,10 @@
+package cdncheck
+
+type Options struct {
+	Cache       bool
+	IPInfoToken string
+}
+
+func (options *Options) HasAuthInfo() bool {
+	return options.IPInfoToken != ""
+}

ranges.go

@@ -10,6 +10,7 @@ import (
 var cidrRegex = regexp.MustCompile(`[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\/[0-9]{1,3}`)
 
 type scraperFunc func(httpClient *http.Client) ([]string, error)
+type scraperWithOptionsFunc func(httpClient *http.Client, options *Options) ([]string, error)
 
 // scrapeAzure scrapes Microsoft Azure firewall's CIDR ranges from their datacenter
 func scrapeAzure(httpClient *http.Client) ([]string, error) {
@@ -90,8 +91,12 @@ func scrapeIncapsula(httpClient *http.Client) ([]string, error) {
 }
 
 // scrapeAkamai scrapes akamai firewall's CIDR ranges from ipinfo
-func scrapeAkamai(httpClient *http.Client) ([]string, error) {
-	resp, err := httpClient.Get("https://ipinfo.io/AS12222")
+func scrapeAkamai(httpClient *http.Client, options *Options) ([]string, error) {
+	req, err := makeReqWithAuth(http.MethodGet, "https://ipinfo.io/AS12222", "Authorization", "Bearer "+options.IPInfoToken)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := httpClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -108,8 +113,12 @@ func scrapeAkamai(httpClient *http.Client) ([]string, error) {
 }
 
 // scrapeSucuri scrapes sucuri firewall's CIDR ranges from ipinfo
-func scrapeSucuri(httpClient *http.Client) ([]string, error) {
-	resp, err := httpClient.Get("https://ipinfo.io/AS30148")
+func scrapeSucuri(httpClient *http.Client, options *Options) ([]string, error) {
+	req, err := makeReqWithAuth(http.MethodGet, "https://ipinfo.io/AS30148", "Authorization", "Bearer "+options.IPInfoToken)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := httpClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -143,8 +152,12 @@ func scrapeFastly(httpClient *http.Client) ([]string, error) {
 }
 
 // scrapeLeaseweb scrapes leaseweb firewall's CIDR ranges from ipinfo
-func scrapeLeaseweb(httpClient *http.Client) ([]string, error) {
-	resp, err := httpClient.Get("https://ipinfo.io/AS60626")
+func scrapeLeaseweb(httpClient *http.Client, options *Options) ([]string, error) {
+	req, err := makeReqWithAuth(http.MethodGet, "https://ipinfo.io/AS60626", "Authorization", "Bearer "+options.IPInfoToken)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := httpClient.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -176,3 +189,12 @@ func scrapeProjectDiscovery(httpClient *http.Client) ([]string, error) {
 	cidrs := cidrRegex.FindAllString(body, -1)
 	return cidrs, nil
 }
+
+func makeReqWithAuth(method, URL, headerName, bearerValue string) (*http.Request, error) {
+	req, err := http.NewRequest(http.MethodGet, URL, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Add(headerName, "Bearer "+bearerValue)
+	return req, nil
+}

ranges_test.go

@@ -23,39 +23,33 @@ func TestScrapeRanges(t *testing.T) {
 	}
 
 	t.Run("azure", func(t *testing.T) {
-		_, err := scrapeAzure(httpClient)
+		ips, err := scrapeAzure(httpClient)
 		require.Nil(t, err, "Could not scrape azure")
+		require.Positive(t, len(ips), "Empty ip list")
 	})
 	t.Run("cloudfront", func(t *testing.T) {
-		_, err := scrapeCloudFront(httpClient)
+		ips, err := scrapeCloudFront(httpClient)
 		require.Nil(t, err, "Could not scrape cloudfront")
+		require.Positive(t, len(ips), "Empty ip list")
 	})
 	t.Run("cloudflare", func(t *testing.T) {
-		_, err := scrapeCloudflare(httpClient)
+		ips, err := scrapeCloudflare(httpClient)
 		require.Nil(t, err, "Could not scrape cloudflare")
+		require.Positive(t, len(ips), "Empty ip list")
 	})
 	t.Run("incapsula", func(t *testing.T) {
-		_, err := scrapeIncapsula(httpClient)
+		ips, err := scrapeIncapsula(httpClient)
 		require.Nil(t, err, "Could not scrape incapsula")
-	})
-	t.Run("akamai", func(t *testing.T) {
-		_, err := scrapeAkamai(httpClient)
-		require.Nil(t, err, "Could not scrape akamai")
-	})
-	t.Run("sucuri", func(t *testing.T) {
-		_, err := scrapeSucuri(httpClient)
-		require.Nil(t, err, "Could not scrape sucuri")
-	})
-	t.Run("leaseweb", func(t *testing.T) {
-		_, err := scrapeLeaseweb(httpClient)
-		require.Nil(t, err, "Could not scrape leaseweb")
+		require.Positive(t, len(ips), "Empty ip list")
 	})
 	t.Run("fastly", func(t *testing.T) {
-		_, err := scrapeFastly(httpClient)
+		ips, err := scrapeFastly(httpClient)
 		require.Nil(t, err, "Could not scrape fastly")
+		require.Positive(t, len(ips), "Empty ip list")
 	})
 	t.Run("projectdiscovery", func(t *testing.T) {
-		_, err := scrapeProjectDiscovery(httpClient)
+		ips, err := scrapeProjectDiscovery(httpClient)
 		require.Nil(t, err, "Could not scrape projectdiscovery")
+		require.Positive(t, len(ips), "Empty ip list")
 	})
 }