ExploreDNS/internal/output/runner.go

package output

import (
	"context"
	"fmt"
	"net"

	"gitea.hansenits.com.au/hits/ExploreDNS/internal/fingerprint"
	"gitea.hansenits.com.au/hits/ExploreDNS/internal/traverse"
)

// RunTraversal drives one traversal and streams its output through the
// formatter. It returns the synthetic root referral whose Stats aggregate
// every leaf outcome.
func RunTraversal(ctx context.Context, traverser *traverse.Traverser, cfg *Config, formatter Formatter, domain string) (*traverse.Referral, error) {
	if traverser == nil {
		return nil, fmt.Errorf("traverser is required")
	}
	if cfg == nil {
		cfg = DefaultConfig()
	}
	if formatter == nil {
		formatter = NewFormatter(cfg, nil)
	}

	traverser.SetHooks(AttachHooks(cfg, formatter))

	// Discover the roots up front so the header can report the initial root;
	// Run reuses the memoised discovery.
	roots, err := traverser.Roots(ctx)
	if err != nil {
		return nil, err
	}
	if err := formatter.WriteHeader(roots); err != nil {
		return nil, err
	}

	root, err := traverser.Run(ctx, domain)
	if err != nil {
		return root, err
	}

	servers := traverser.ServersEncountered()

	// Fingerprint servers when both ShowVersions and ShowServers are enabled.
	// Gating on ShowServers avoids unnecessary network calls when versions
	// would not be displayed anyway.
	if cfg.ShowVersions && cfg.ShowServers {
		cfg.Fingerprints = fingerprint.New().FingerprintAll(ctx, collectUniqueServerIPs(servers))
	}

	if err := formatter.WriteSummary(root, servers); err != nil {
		return root, err
	}
	if err := formatter.Flush(); err != nil {
		return root, err
	}

	return root, nil
}

// collectUniqueServerIPs returns the set of unique server IPs encountered.
func collectUniqueServerIPs(servers map[string][]string) []net.IP {
	seen := make(map[string]bool)
	var ips []net.IP
	for _, addrs := range servers {
		for _, addr := range addrs {
			if seen[addr] {
				continue
			}
			seen[addr] = true
			if ip := net.ParseIP(addr); ip != nil {
				ips = append(ips, ip)
			}
		}
	}
	return ips
}