支持多DNS服务器和多线程

Author: Anion3r

src/main.go

@@ -8,17 +8,26 @@ import (
 	_ "github.com/miekg/dns"
 	"io"
 	"log"
+	"net"
 	"os"
 	"strings"
+	"sync"
 )
 
 const DEBUG = false
 const A uint16 = 0x01
+const MaxThread int = 100
+
+var tmpChain = make(chan struct{}, MaxThread)
+var waitGroup sync.WaitGroup
 
 func GetAName(DNSServer, domain string) bool {
+	defer waitGroup.Done()
 	msg := new(dns.Msg)
 	msg.SetQuestion(dns.Fqdn(domain), dns.TypeA)
+	tmpChain <- struct{}{}
 	in, err := dns.Exchange(msg, DNSServer+":53")
+	<-tmpChain
 	if err != nil {
 		if DEBUG {
 			log.Printf("[x] %s", err)
@@ -38,7 +47,7 @@ func usage() {
 	fmt.Fprintf(os.Stderr,
 		`GoDigDomain: A Domain Name Burst Tool
 
-Usage: gdd -dn Domain [-ds DNSServer] [-dt DictFile]
+Usage: gdd -dn Domain [-ds DNSServer(s)] [-dt DictFile]
 
 Options:
 `)
@@ -53,11 +62,23 @@ func isFileExist(path string) bool {
 	return true
 }
 
+func createDNSServerList(DNSServers string) []string {
+	dnsServers := strings.Split(DNSServers, ",")
+	ret := make([]string, len(dnsServers))
+	for i := 0; i < len(dnsServers); i++ {
+		if net.ParseIP(dnsServers[i]) == nil {
+			continue
+		}
+		ret[i] = dnsServers[i]
+	}
+	return ret
+}
+
 func main() {
 	log.SetFlags(log.Lshortfile)
 	flag.Usage = usage
 
-	dnsServer := flag.String("ds", "114.114.114.114", "DNS服务器")
+	dnsServer := flag.String("ds", "114.114.114.114", "DNS服务器(多个服务器用逗号[,]隔开)")
 	domain := flag.String("dn", "", "域名")
 	dict := flag.String("dt", "./dict.txt", "域名字典")
 	flag.Parse()
@@ -82,6 +103,8 @@ func main() {
 	defer f.Close()
 
 	lines := bufio.NewReader(f)
+	dnsServers := createDNSServerList(*dnsServer)
+
 	for {
 		line, _, err := lines.ReadLine()
 		if err != nil {
@@ -96,7 +119,10 @@ func main() {
 			return
 		}
 		domainPrefix := string(line)
-		GetAName(*dnsServer, fmt.Sprintf("%s.%s", strings.TrimPrefix(domainPrefix, "\n\r"), *domain))
+		for i := 0; i < len(dnsServers); i++ {
+			waitGroup.Add(1)
+			go GetAName(dnsServers[i], fmt.Sprintf("%s.%s", strings.TrimPrefix(domainPrefix, "\n\r"), *domain))
+		}
 	}
-
+	waitGroup.Wait()
 }