Improved DNS handling, pull servers from trickest when none supplied, added CNAME handling, updated preview picture

ptrstream.go

@@ -7,6 +7,7 @@ import (
 	"flag"
 	"fmt"
 	"net"
+	"net/http"
 	"os"
 	"regexp"
 	"strconv"
@@ -19,6 +20,8 @@ import (
 	"github.com/rivo/tview"
 )
 
+const defaultResolversURL = "https://raw.githubusercontent.com/trickest/resolvers/refs/heads/main/resolvers.txt"
+
 type Config struct {
 	concurrency   int
 	timeout       time.Duration
@@ -28,6 +31,8 @@ type Config struct {
 	debug         bool
 	outputFile    *os.File
 	mu            sync.Mutex
+	lastDNSUpdate time.Time
+	updateMu      sync.Mutex
 }
 
 type Stats struct {
@@ -37,6 +42,7 @@ type Stats struct {
 	lastCheckTime time.Time
 	success       uint64
 	failed        uint64
+	cnames        uint64
 	speedHistory  []float64
 	mu            sync.Mutex
 }
@@ -53,7 +59,15 @@ func (s *Stats) incrementFailed() {
 	atomic.AddUint64(&s.failed, 1)
 }
 
+func (s *Stats) incrementCNAME() {
+	atomic.AddUint64(&s.cnames, 1)
+}
+
 func (c *Config) getNextServer() string {
+	if err := c.updateDNSServers(); err != nil {
+		fmt.Printf("Failed to update DNS servers: %v\n", err)
+	}
+
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -66,40 +80,81 @@ func (c *Config) getNextServer() string {
 	return server
 }
 
-func loadDNSServers(filename string) ([]string, error) {
+func fetchDefaultResolvers() ([]string, error) {
-	if filename == "" {
+	resp, err := http.Get(defaultResolversURL)
-		return nil, nil
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch default resolvers: %v", err)
+	}
+	defer resp.Body.Close()
+
+	var resolvers []string
+	scanner := bufio.NewScanner(resp.Body)
+	for scanner.Scan() {
+		resolver := strings.TrimSpace(scanner.Text())
+		if resolver != "" {
+			resolvers = append(resolvers, resolver)
+		}
 	}
 
-	file, err := os.Open(filename)
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("error reading default resolvers: %v", err)
+	}
+
+	return resolvers, nil
+}
+
+func loadDNSServers(dnsFile string) ([]string, error) {
+	if dnsFile == "" {
+		resolvers, err := fetchDefaultResolvers()
 		if err != nil {
 			return nil, err
 		}
+		if len(resolvers) == 0 {
+			return nil, fmt.Errorf("no default resolvers found")
+		}
+		return resolvers, nil
+	}
+
+	file, err := os.Open(dnsFile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open DNS servers file: %v", err)
+	}
 	defer file.Close()
 
 	var servers []string
 	scanner := bufio.NewScanner(file)
 	for scanner.Scan() {
 		server := strings.TrimSpace(scanner.Text())
-		if server != "" && !strings.HasPrefix(server, "#") {
+		if server != "" {
-			if !strings.Contains(server, ":") {
-				server += ":53"
-			}
 			servers = append(servers, server)
 		}
 	}
 
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("error reading DNS servers file: %v", err)
+	}
+
 	if len(servers) == 0 {
-		return nil, fmt.Errorf("no valid DNS servers found in file")
+		return nil, fmt.Errorf("no DNS servers found in file")
 	}
 
-	return servers, scanner.Err()
+	return servers, nil
 }
 
-func lookupWithRetry(ip string, cfg *Config) ([]string, string, error) {
+type DNSResponse struct {
+	Names      []string
+	Server     string
+	RecordType string // "PTR" or "CNAME"
+	Target     string // For CNAME records, stores the target
+}
+
+func lookupWithRetry(ip string, cfg *Config) (DNSResponse, error) {
+	var lastErr error
+
+	for i := 0; i < cfg.retries; i++ {
 		server := cfg.getNextServer()
 		if server == "" {
-		return nil, "", fmt.Errorf("no DNS servers available")
+			return DNSResponse{}, fmt.Errorf("no DNS servers available")
 		}
 
 		r := &net.Resolver{
@@ -112,32 +167,43 @@ func lookupWithRetry(ip string, cfg *Config) ([]string, string, error) {
 			},
 		}
 
-	for i := 0; i < cfg.retries; i++ {
 		ctx, cancel := context.WithTimeout(context.Background(), cfg.timeout)
 		names, err := r.LookupAddr(ctx, ip)
 		cancel()
 
 		if err == nil {
-			return names, server, nil
+			logServer := server
+			if idx := strings.Index(server, ":"); idx != -1 {
+				logServer = server[:idx]
 			}
 
-		if i < cfg.retries-1 {
+			// Check if any of the names is a CNAME
-			server = cfg.getNextServer()
+			for _, name := range names {
-			if server == "" {
+				ctx, cancel := context.WithTimeout(context.Background(), cfg.timeout)
-				return nil, "", fmt.Errorf("no more DNS servers available")
+				cname, err := r.LookupCNAME(ctx, strings.TrimSuffix(name, "."))
-			}
+				cancel()
-			r = &net.Resolver{
+
-				PreferGo: true,
+				if err == nil && cname != name {
-				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+					return DNSResponse{
-					d := net.Dialer{
+						Names:      names,
-						Timeout: cfg.timeout,
+						Server:     logServer,
-					}
+						RecordType: "CNAME",
-					return d.DialContext(ctx, "udp", server)
+						Target:     strings.TrimSuffix(cname, "."),
-				},
+					}, nil
 				}
 			}
+
+			return DNSResponse{
+				Names:      names,
+				Server:     logServer,
+				RecordType: "PTR",
+			}, nil
 		}
-	return nil, "", fmt.Errorf("lookup failed after %d retries", cfg.retries)
+
+		lastErr = err
+	}
+
+	return DNSResponse{}, lastErr
 }
 
 func reverse(ss []string) []string {
@@ -201,9 +267,16 @@ func colorizeIPInPtr(ptr, ip string) string {
 	}
 
 	finalResult := result.String()
-	finalResult = strings.ReplaceAll(finalResult, ".in-addr.arpa", ".[blue]in-addr.arpa")
+
-	finalResult = strings.ReplaceAll(finalResult, ".gov", ".[red]gov")
+	if strings.HasSuffix(finalResult, ".in-addr.arpa") {
-	finalResult = strings.ReplaceAll(finalResult, ".mil", ".[red]mil")
+		finalResult = finalResult[:len(finalResult)-13] + ".[blue]in-addr.arpa"
+	}
+	if strings.HasSuffix(finalResult, ".gov") {
+		finalResult = finalResult[:len(finalResult)-4] + ".[red]gov"
+	}
+	if strings.HasSuffix(finalResult, ".mil") {
+		finalResult = finalResult[:len(finalResult)-4] + ".[red]mil"
+	}
 
 	return finalResult
 }
@@ -219,18 +292,17 @@ const maxBufferLines = 1000
 func worker(jobs <-chan string, wg *sync.WaitGroup, cfg *Config, stats *Stats, textView *tview.TextView, app *tview.Application) {
 	defer wg.Done()
 	for ip := range jobs {
-		var names []string
-		var server string
-		var err error
 		timestamp := time.Now()
+		var response DNSResponse
+		var err error
 
 		if len(cfg.dnsServers) > 0 {
-			names, server, err = lookupWithRetry(ip, cfg)
+			response, err = lookupWithRetry(ip, cfg)
-			if idx := strings.Index(server, ":"); idx != -1 {
-				server = server[:idx]
-			}
 		} else {
-			names, err = net.LookupAddr(ip)
+			names, err := net.LookupAddr(ip)
+			if err == nil {
+				response = DNSResponse{Names: names, RecordType: "PTR"}
+			}
 		}
 
 		stats.increment()
@@ -255,7 +327,7 @@ func worker(jobs <-chan string, wg *sync.WaitGroup, cfg *Config, stats *Stats, t
 			continue
 		}
 
-		if len(names) == 0 {
+		if len(response.Names) == 0 {
 			stats.incrementFailed()
 			if cfg.debug {
 				timestamp := time.Now().Format("2006-01-02 15:04:05")
@@ -273,7 +345,7 @@ func worker(jobs <-chan string, wg *sync.WaitGroup, cfg *Config, stats *Stats, t
 		stats.incrementSuccess()
 
 		ptr := ""
-		for _, name := range names {
+		for _, name := range response.Names {
 			if cleaned := strings.TrimSpace(strings.TrimSuffix(name, ".")); cleaned != "" {
 				ptr = cleaned
 				break
@@ -284,21 +356,29 @@ func worker(jobs <-chan string, wg *sync.WaitGroup, cfg *Config, stats *Stats, t
 			continue
 		}
 
-		writeNDJSON(cfg, timestamp, ip, server, ptr)
+		writeNDJSON(cfg, timestamp, ip, response.Server, ptr, response.RecordType, response.Target)
 
 		timeStr := time.Now().Format("2006-01-02 15:04:05")
+		recordTypeColor := "[blue] PTR [-]"
+		if response.RecordType == "CNAME" {
+			stats.incrementCNAME()
+			recordTypeColor = "[fuchsia]CNAME[-]"
+			ptr = fmt.Sprintf("%s -> %s", ptr, response.Target)
+		}
 
 		var line string
 		if len(cfg.dnsServers) > 0 {
-			line = fmt.Sprintf("[gray]%s[-] [purple]%15s[-] [gray]│[-] [yellow]%15s[-] [gray]│[-] %s\n",
+			line = fmt.Sprintf("[gray]%s [gray]│[-] [purple]%15s[-] [gray]│[-] [yellow]%-15s[-] [gray]│[-] %-5s [gray]│[-] %s\n",
 				timeStr,
 				ip,
-				server,
+				response.Server,
+				recordTypeColor,
 				colorizeIPInPtr(ptr, ip))
 		} else {
-			line = fmt.Sprintf("[gray]%s[-] [purple]%15s[-] [gray]│[-] %s\n",
+			line = fmt.Sprintf("[gray]%s [gray]│[-] [purple]%15s[-] [gray]│[-] %-5s [gray]│[-] %s\n",
 				timeStr,
 				ip,
+				recordTypeColor,
 				colorizeIPInPtr(ptr, ip))
 		}
 
@@ -343,6 +423,38 @@ func parseShardArg(shard string) (int, int, error) {
 	return shardNum, totalShards, nil
 }
 
+func (c *Config) updateDNSServers() error {
+	c.updateMu.Lock()
+	defer c.updateMu.Unlock()
+
+	if time.Since(c.lastDNSUpdate) < 24*time.Hour {
+		return nil
+	}
+
+	resolvers, err := fetchDefaultResolvers()
+	if err != nil {
+		return err
+	}
+
+	if len(resolvers) == 0 {
+		return fmt.Errorf("no resolvers found in update")
+	}
+
+	for i, server := range resolvers {
+		if !strings.Contains(server, ":") {
+			resolvers[i] = server + ":53"
+		}
+	}
+
+	c.mu.Lock()
+	c.dnsServers = resolvers
+	c.serverIndex = 0
+	c.lastDNSUpdate = time.Now()
+	c.mu.Unlock()
+
+	return nil
+}
+
 func main() {
 	concurrency := flag.Int("c", 100, "Concurrency level")
 	timeout := flag.Duration("t", 2*time.Second, "Timeout for DNS queries")
@@ -364,11 +476,25 @@ func main() {
 		*seed = time.Now().UnixNano()
 	}
 
+	servers, err := loadDNSServers(*dnsFile)
+	if err != nil {
+		fmt.Printf("Error loading DNS servers: %v\n", err)
+		return
+	}
+
+	for i, server := range servers {
+		if !strings.Contains(server, ":") {
+			servers[i] = server + ":53"
+		}
+	}
+
 	cfg := &Config{
 		concurrency:   *concurrency,
 		timeout:       *timeout,
 		retries:       *retries,
 		debug:         *debug,
+		dnsServers:    servers,
+		lastDNSUpdate: time.Now(),
 	}
 
 	if *outputPath != "" {
@@ -381,16 +507,6 @@ func main() {
 		defer f.Close()
 	}
 
-	if *dnsFile != "" {
-		servers, err := loadDNSServers(*dnsFile)
-		if err != nil {
-			fmt.Printf("Error loading DNS servers: %v\n", err)
-			return
-		}
-		cfg.dnsServers = servers
-		fmt.Printf("Loaded %d DNS servers\n", len(servers))
-	}
-
 	app := tview.NewApplication()
 
 	textView := tview.NewTextView().
@@ -456,10 +572,11 @@ func main() {
 						return
 					}
 
-					statsText := fmt.Sprintf(" [aqua]Count:[:-] [white]%s [gray]│[-] [aqua]Progress:[:-] [darkgray]%7.2f%%[-] [gray]│[-] [aqua]Rate:[:-] %s [gray]│[-] [aqua]Successful:[:-] [green]✓%s [-][darkgray](%5.1f%%)[-] [gray]│[-] [aqua]Failed:[:-] [red]✗%s [-][darkgray](%5.1f%%)[-] ",
+					statsText := fmt.Sprintf(" [aqua]Count:[:-] [white]%s [gray]│[-] [aqua]Progress:[:-] [darkgray]%7.2f%%[-] [gray]│[-] [aqua]Rate:[:-] %s [gray]│[-] [aqua]CNAMEs:[:-] [yellow]%s[-] [gray]│[-] [aqua]Successful:[:-] [green]✓%s [-][darkgray](%5.1f%%)[-] [gray]│[-] [aqua]Failed:[:-] [red]✗%s [-][darkgray](%5.1f%%)[-] ",
 						formatNumber(processed),
 						percent,
 						colorizeSpeed(avgSpeed),
+						formatNumber(atomic.LoadUint64(&stats.cnames)),
 						formatNumber(success),
 						float64(success)/float64(processed)*100,
 						formatNumber(failed),
@@ -568,7 +685,7 @@ func visibleLength(s string) int {
 	return len(noColors)
 }
 
-func writeNDJSON(cfg *Config, timestamp time.Time, ip, server, ptr string) {
+func writeNDJSON(cfg *Config, timestamp time.Time, ip, server, ptr, recordType, target string) {
 	if cfg.outputFile == nil {
 		return
 	}
@@ -578,11 +695,15 @@ func writeNDJSON(cfg *Config, timestamp time.Time, ip, server, ptr string) {
 		IPAddr     string `json:"ip_addr"`
 		DNSServer  string `json:"dns_server"`
 		PTRRecord  string `json:"ptr_record"`
+		RecordType string `json:"record_type"`
+		Target     string `json:"target,omitempty"`
 	}{
 		Timestamp:  timestamp.Format(time.RFC3339),
 		IPAddr:     ip,
 		DNSServer:  server,
 		PTRRecord:  ptr,
+		RecordType: recordType,
+		Target:     target,
 	}
 
 	if data, err := json.Marshal(record); err == nil {