Adding hostsfile support

client.go

@@ -13,7 +13,9 @@ import (
 	"time"
 
 	"github.com/miekg/dns"
+	"github.com/projectdiscovery/iputil"
 	"github.com/projectdiscovery/retryabledns/doh"
+	"github.com/projectdiscovery/retryabledns/hostsfile"
 )
 
 func init() {
@@ -28,6 +30,7 @@ type Client struct {
 	TCPFallback  bool
 	tcpClient    *dns.Client
 	dohClient    *doh.Client
+	knownHosts   map[string][]string
 }
 
 // New creates a new dns client
@@ -38,11 +41,16 @@ func New(baseResolvers []string, maxRetries int) *Client {
 // New creates a new dns client with options
 func NewWithOptions(options Options) *Client {
 	parsedBaseResolvers := parseResolvers(deduplicate(options.BaseResolvers))
+	var knownHosts map[string][]string
+	if options.Hostsfile {
+		knownHosts, _ = hostsfile.ParseDefault()
+	}
 	client := Client{
-		options:   options,
-		resolvers: parsedBaseResolvers,
-		tcpClient: &dns.Client{Net: TCP.String(), Timeout: options.Timeout},
-		dohClient: doh.New(),
+		options:    options,
+		resolvers:  parsedBaseResolvers,
+		tcpClient:  &dns.Client{Net: TCP.String(), Timeout: options.Timeout},
+		dohClient:  doh.New(),
+		knownHosts: knownHosts,
 	}
 	return &client
 }
@@ -158,6 +166,19 @@ func (c *Client) QueryMultiple(host string, requestTypes []uint16) (*DNSData, er
 		err     error
 	)
 
+	// integrate data with known hosts in case
+	if c.options.Hostsfile {
+		if ips, ok := c.knownHosts[host]; ok {
+			for _, ip := range ips {
+				if iputil.IsIPv4(ip) {
+					dnsdata.A = append(dnsdata.A, ip)
+				} else if iputil.IsIPv6(ip) {
+					dnsdata.AAAA = append(dnsdata.AAAA, ip)
+				}
+			}
+		}
+	}
+
 	msg := &dns.Msg{}
 	msg.Id = dns.Id()
 	msg.RecursionDesired = true
@@ -237,6 +258,7 @@ func (c *Client) QueryMultiple(host string, requestTypes []uint16) (*DNSData, er
 			}
 		}
 	}
+
 	return &dnsdata, err
 }
 

go.mod

@@ -5,6 +5,7 @@ go 1.14
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/miekg/dns v1.1.29
+	github.com/projectdiscovery/fileutil v0.0.0-20210926202739-6050d0acf73c
 	github.com/projectdiscovery/retryablehttp-go v1.0.2
 	github.com/projectdiscovery/stringsutil v0.0.0-20210823090203-2f5f137e8e1d
 	github.com/stretchr/testify v1.7.0

go.sum

@@ -1,10 +1,14 @@
 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/karrick/godirwalk v1.16.1 h1:DynhcF+bztK8gooS0+NDJFrdNZjJ3gzVzC545UNA9iw=
+github.com/karrick/godirwalk v1.16.1/go.mod h1:j4mkqPuvaLI8mp1DroR3P6ad7cyYd4c1qeJ3RV7ULlk=
 github.com/miekg/dns v1.1.29 h1:xHBEhR+t5RzcFJjBLJlax2daXOrTYtr9z4WdKEfWFzg=
 github.com/miekg/dns v1.1.29/go.mod h1:KNUDUusw/aVsxyTYZM1oqvCicbwhgbNgztCETuNZ7xM=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/projectdiscovery/fileutil v0.0.0-20210926202739-6050d0acf73c h1:KDmCXhLLnS/Gc1VDyTxxamRzc8OmHCm1X+f8WQoaTRs=
+github.com/projectdiscovery/fileutil v0.0.0-20210926202739-6050d0acf73c/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
 github.com/projectdiscovery/retryablehttp-go v1.0.2 h1:LV1/KAQU+yeWhNVlvveaYFsjBYRwXlNEq0PvrezMV0U=
 github.com/projectdiscovery/retryablehttp-go v1.0.2/go.mod h1:dx//aY9V247qHdsRf0vdWHTBZuBQ2vm6Dq5dagxrDYI=
 github.com/projectdiscovery/stringsutil v0.0.0-20210823090203-2f5f137e8e1d h1:lrdpJCBOvRrTnm44Ov7O3tLd3oOWhCvVUhTKkWwibq4=

hostsfile/hostsfile.go

@@ -0,0 +1,56 @@
+package hostsfile
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"runtime"
+	"strings"
+
+	"github.com/projectdiscovery/fileutil"
+	"github.com/projectdiscovery/stringsutil"
+)
+
+func Path() string {
+	if runtime.GOOS == "windows" {
+		return fmt.Sprintf(`%s\System32\Drivers\etc\hosts`, os.Getenv("SystemRoot"))
+	}
+	return "/etc/hosts"
+}
+
+func ParseDefault() (map[string][]string, error) {
+	return Parse(Path())
+}
+
+func Parse(p string) (map[string][]string, error) {
+	if !fileutil.FileExists(p) {
+		return nil, errors.New("hosts file doesn't exist")
+	}
+
+	hostsFileCh, err := fileutil.ReadFile(p)
+	if err != nil {
+		return nil, err
+	}
+
+	items := make(map[string][]string)
+	for line := range hostsFileCh {
+		line = strings.TrimSpace(line)
+		// skip comments and empty lines
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+
+		// discard comment part
+		if strings.Contains(line, "#") {
+			line = stringsutil.Before(line, "#")
+		}
+		tokens := strings.Fields(line)
+		if len(tokens) > 1 {
+			ip := tokens[0]
+			for _, hostname := range tokens[1:] {
+				items[hostname] = append(items[hostname], ip)
+			}
+		}
+	}
+	return items, nil
+}

options.go

@@ -6,4 +6,5 @@ type Options struct {
 	BaseResolvers []string
 	MaxRetries    int
 	Timeout       time.Duration
+	Hostsfile     bool
 }