Add support for custom name servers

This allows you to override the default name server as specified in
/etc/resolv.conf.

.gitignore

@@ -0,0 +1,2 @@
+s3enum
+*.txt

bucket_checker.go

@@ -4,21 +4,30 @@ import (
 	"errors"
 	"fmt"
 	"github.com/miekg/dns"
+	"net"
 	"strings"
 )
 
 type Resolver interface {
 	IsBucket(string) bool
 }
 
-func NewS3Resolver() *S3Resolver {
+func NewS3Resolver(ns string) (*S3Resolver, error) {
+	config, err := getConfig(ns)
+
+	if err != nil {
+		return nil, err
+	}
+
 	return &S3Resolver{
 		dnsClient: dns.Client{},
-	}
+		config:    *config,
+	}, nil
 }
 
 type S3Resolver struct {
 	dnsClient dns.Client
+	config    dns.ClientConfig
 }
 
 const s3host = "s3.amazonaws.com"
@@ -34,12 +43,34 @@ func (s *S3Resolver) IsBucket(name string) bool {
 	return false
 }
 
+func getConfig(nameserver string) (*dns.ClientConfig, error) {
+	if nameserver != "" {
+		addr := net.ParseIP(nameserver)
+		if addr != nil {
+			return &dns.ClientConfig{
+				Servers: []string{addr.String()},
+				Port:    "53",
+			}, nil
+		} else {
+			return nil, errors.New("invalid ip addr")
+		}
+	} else {
+		config, err := dns.ClientConfigFromFile("/etc/resolv.conf")
+
+		if err != nil {
+			return nil, errors.New("could not read local resolver config")
+		}
+
+		return config, nil
+	}
+}
+
 func (s *S3Resolver) resolveCNAME(name string) (string, error) {
 	msg := dns.Msg{}
 	msg.SetQuestion(name, dns.TypeCNAME)
 
-	// TODO: Allow the name server to be set by the user.
-	r, _, err := s.dnsClient.Exchange(&msg, "8.8.8.8:53")
+	addr := net.JoinHostPort(s.config.Servers[0], s.config.Port)
+	r, _, err := s.dnsClient.Exchange(&msg, addr)
 
 	if err != nil {
 		return "", errors.New("probably a timeout")

main.go

@@ -11,21 +11,23 @@ var (
 	names              []string
 	wordListFile       string
 	preAndSuffixesFile string
+	nameserver         string
 )
 
 const version = "0.0.1"
 const usage = `s3enum
 
 Usage:
-  s3enum --wordlist wl.txt --suffixlist sl.txt [--threads 2] <name>...
+  s3enum --wordlist wl.txt --suffixlist sl.txt [--threads 2] [--nameserver 1.1.1.1] <name>...
   s3enum -h | --help
   s3enum --version
 
 Options:
-  --wordlist <path>    Path to the word list.
-  --suffixlist <path>  Path to the word list.
-  --threads <threads>  Number of threads [default: 10].
-  -h --help            Show this screen.`
+  --wordlist <path>             Path to the word list.
+  --suffixlist <path>           Path to the word list.
+  --threads <threads>           Number of threads [default: 10].
+  -n --nameserver <nameserver>  Use specific nameserver.
+  -h --help                     Show this screen.`
 
 func main() {
 	opts, err := docopt.ParseDoc(usage)
@@ -43,13 +45,23 @@ func main() {
 	wordListFile = opts["--wordlist"].(string)
 	threads, _ = opts.Int("--threads")
 
+	if opts["--nameserver"] == nil {
+		nameserver = ""
+	} else {
+		nameserver = opts["--nameserver"].(string)
+	}
+
 	wordChannel := make(chan string)
 	wordDone := make(chan bool)
 
 	resultChannel := make(chan string)
 	resultDone := make(chan bool)
 
-	resolver := NewS3Resolver()
+	resolver, err := NewS3Resolver(nameserver)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Could not initialize DNS resolver: %v\n", err)
+		os.Exit(1)
+	}
 
 	consumer, err := NewConsumer(resolver, wordChannel, resultChannel, wordDone)
 	if err != nil {