Skip to content

Commit

Permalink
Merge pull request #37 from projectdiscovery/add-hostsfile-support
Browse files Browse the repository at this point in the history
Adding hostsfile support
  • Loading branch information
Mzack9999 committed Oct 11, 2021
2 parents 76c5b76 + f4a76fa commit 9d3cec3
Show file tree
Hide file tree
Showing 5 changed files with 88 additions and 4 deletions.
30 changes: 26 additions & 4 deletions client.go
Expand Up @@ -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() {
Expand All @@ -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
Expand All @@ -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
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -237,6 +258,7 @@ func (c *Client) QueryMultiple(host string, requestTypes []uint16) (*DNSData, er
}
}
}

return &dnsdata, err
}

Expand Down
1 change: 1 addition & 0 deletions go.mod
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions 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=
Expand Down
56 changes: 56 additions & 0 deletions 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
}
1 change: 1 addition & 0 deletions options.go
Expand Up @@ -6,4 +6,5 @@ type Options struct {
BaseResolvers []string
MaxRetries int
Timeout time.Duration
Hostsfile bool
}

0 comments on commit 9d3cec3

Please sign in to comment.