ExploreDNS/internal/output/formatter_test.go

package output

import (
	"bytes"
	"context"
	"encoding/json"
	"net"
	"strings"
	"testing"
	"time"

	idns "gitea.hansenits.com.au/hits/ExploreDNS/internal/dns"
	"gitea.hansenits.com.au/hits/ExploreDNS/internal/traverse"
	"github.com/miekg/dns"
)

// mockDelegation wires root → com → example.com (2 NS, one glueless answer
// path) through the single injected exchange.
func mockDelegation() idns.ExchangeFunc {
	responses := map[string]*dns.Msg{}
	set := func(server, qname string, msg *dns.Msg) {
		responses[server+"/"+dns.Fqdn(qname)] = msg
	}
	a := func(name, ip string) dns.RR {
		return &dns.A{
			Hdr: dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300},
			A:   net.ParseIP(ip).To4(),
		}
	}
	ns := func(zone, target string) dns.RR {
		return &dns.NS{
			Hdr: dns.RR_Header{Name: dns.Fqdn(zone), Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 300},
			Ns:  dns.Fqdn(target),
		}
	}

	rootMsg := new(dns.Msg)
	rootMsg.Ns = []dns.RR{ns("com", "a.gtld-servers.net")}
	rootMsg.Extra = []dns.RR{a("a.gtld-servers.net", "192.5.6.30")}
	set("198.41.0.4", "www.example.com", rootMsg)

	comMsg := new(dns.Msg)
	comMsg.Ns = []dns.RR{ns("example.com", "ns1.example.com"), ns("example.com", "ns2.example.com")}
	comMsg.Extra = []dns.RR{a("ns1.example.com", "1.1.1.1"), a("ns2.example.com", "2.2.2.2")}
	set("192.5.6.30", "www.example.com", comMsg)

	answer := new(dns.Msg)
	answer.Answer = []dns.RR{a("www.example.com", "9.9.9.9")}
	set("1.1.1.1", "www.example.com", answer)
	set("2.2.2.2", "www.example.com", answer)

	return func(_ context.Context, server string, msg *dns.Msg, _ bool) (*dns.Msg, error) {
		host := server
		if h, _, err := net.SplitHostPort(server); err == nil {
			host = h
		}
		resp, ok := responses[host+"/"+msg.Question[0].Name]
		if !ok {
			return nil, &net.DNSError{Err: "no mock", Name: msg.Question[0].Name}
		}
		out := resp.Copy()
		out.SetReply(msg)
		out.Answer, out.Ns, out.Extra = resp.Answer, resp.Ns, resp.Extra
		return out, nil
	}
}

func newMockTraverser() *traverse.Traverser {
	tr := traverse.NewTraverser(&traverse.TraverserConfig{
		MaxDepth:  traverse.DefaultMaxDepth,
		QueryType: dns.TypeA,
		RootAddrs: []net.IP{net.ParseIP("198.41.0.4")},
		QueryConfig: &idns.QueryConfig{
			Retries:    1,
			Timeout:    time.Second,
			RetryDelay: time.Millisecond,
		},
	})
	tr.SetExchange(mockDelegation())
	return tr
}

func TestDefaultConfig(t *testing.T) {
	cfg := DefaultConfig()
	if cfg.Format != FormatText || !cfg.ShowProgress || !cfg.ShowResults {
		t.Errorf("unexpected defaults: %+v", cfg)
	}
}

func TestNewFormatterSelectsImplementation(t *testing.T) {
	if _, ok := NewFormatter(&Config{Format: FormatJSON}, &bytes.Buffer{}).(*jsonFormatter); !ok {
		t.Error("FormatJSON should select the JSON formatter")
	}
	if _, ok := NewFormatter(&Config{Format: FormatText}, &bytes.Buffer{}).(*textFormatter); !ok {
		t.Error("FormatText should select the text formatter")
	}
	if NewFormatter(nil, &bytes.Buffer{}) == nil {
		t.Error("nil config should still produce a formatter")
	}
}

func TestRunTraversalTextOutput(t *testing.T) {
	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Domain = "www.example.com"
	cfg.QueryType = "a"
	cfg.ShowServers = true
	cfg.ShowVersions = false // no fingerprint network calls in tests
	cfg.Color = false
	formatter := NewFormatter(cfg, &buf)

	root, err := RunTraversal(context.Background(), newMockTraverser(), cfg, formatter, "www.example.com")
	if err != nil {
		t.Fatalf("RunTraversal: %v", err)
	}
	if root == nil || len(root.Stats) == 0 {
		t.Fatal("expected aggregated stats on the root referral")
	}

	out := buf.String()
	for _, want := range []string{
		"# Using fast mode",
		"Using 198.41.0.4 (198.41.0.4) as initial root",
		"Running query www.example.com type a",
		"1 198.41.0.4 (198.41.0.4)",
		"1.1 a.gtld-servers.net (192.5.6.30)",
		"1.1.1 ns1.example.com (1.1.1.1)",
		"Results:",
		" 50.0%: Answer from ns1.example.com (1.1.1.1)",
		" 50.0%: Answer from ns2.example.com (2.2.2.2)",
		"Summary Results:",
		" 100% answered with www.example.com. 300 IN A 9.9.9.9",
		"The following servers were encountered:",
	} {
		if !strings.Contains(out, want) {
			t.Errorf("output missing %q\n---\n%s", want, out)
		}
	}
}

func TestVerboseProgressFormat(t *testing.T) {
	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Domain = "www.example.com"
	cfg.QueryType = "a"
	cfg.ShowVersions = false
	cfg.Verbose = true
	cfg.Color = false
	formatter := NewFormatter(cfg, &buf)

	if _, err := RunTraversal(context.Background(), newMockTraverser(), cfg, formatter, "www.example.com"); err != nil {
		t.Fatalf("RunTraversal: %v", err)
	}
	out := buf.String()
	// Verbose rows are "<refid> [qname] <server> (<ips>) <bailiwick>"; the
	// root bailiwick renders as "<>".
	for _, want := range []string{
		"1 [www.example.com] 198.41.0.4 (198.41.0.4) <>",
		"1.1 [www.example.com] a.gtld-servers.net (192.5.6.30) <com>",
		"1.1.1 [www.example.com] ns1.example.com (1.1.1.1) <example.com>",
	} {
		if !strings.Contains(out, want) {
			t.Errorf("output missing %q\n---\n%s", want, out)
		}
	}
}

func TestRunTraversalQuietSuppressesHeader(t *testing.T) {
	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Domain = "www.example.com"
	cfg.QueryType = "a"
	cfg.ShowVersions = false
	cfg.Quiet = true
	cfg.Color = false
	formatter := NewFormatter(cfg, &buf)

	if _, err := RunTraversal(context.Background(), newMockTraverser(), cfg, formatter, "www.example.com"); err != nil {
		t.Fatalf("RunTraversal: %v", err)
	}
	out := buf.String()
	for _, banned := range []string{"# Using fast mode", "as initial root", "Running query"} {
		if strings.Contains(out, banned) {
			t.Errorf("quiet output must not contain %q\n---\n%s", banned, out)
		}
	}
	if !strings.Contains(out, "Results:") {
		t.Errorf("quiet must still print results\n---\n%s", out)
	}
}

func TestRunTraversalJSONOutput(t *testing.T) {
	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Format = FormatJSON
	cfg.Domain = "www.example.com"
	cfg.QueryType = "A"
	cfg.ShowVersions = false
	formatter := NewFormatter(cfg, &buf)

	if _, err := RunTraversal(context.Background(), newMockTraverser(), cfg, formatter, "www.example.com"); err != nil {
		t.Fatalf("RunTraversal: %v", err)
	}

	var doc map[string]any
	if err := json.Unmarshal(buf.Bytes(), &doc); err != nil {
		t.Fatalf("invalid JSON: %v\n%s", err, buf.String())
	}
	if doc["domain"] != "www.example.com" {
		t.Errorf("domain = %v", doc["domain"])
	}
	if doc["qtype"] != "A" {
		t.Errorf("qtype = %v", doc["qtype"])
	}
	root, ok := doc["root"].(map[string]any)
	if !ok || root["ip"] != "198.41.0.4" {
		t.Errorf("root = %v", doc["root"])
	}
	for _, banned := range []string{"progress", "resolves"} {
		if _, present := doc[banned]; present {
			t.Errorf("JSON document must not contain %q", banned)
		}
	}
	results, ok := doc["results"].([]any)
	if !ok || len(results) != 2 {
		t.Fatalf("results = %v", doc["results"])
	}
	summary, ok := doc["summary"].(map[string]any)
	if !ok {
		t.Fatalf("summary missing: %v", doc)
	}
	byStatus := summary["by_status"].(map[string]any)
	if prob := byStatus["answered"].(float64); prob < 0.999 || prob > 1.001 {
		t.Errorf("answered summary prob = %v", prob)
	}
}

// mockGluelessDelegation wires root → com → example.com where the single NS
// (ns1.example.net) comes without glue, forcing a resolve subtree that walks
// root → net → answer.
func mockGluelessDelegation() idns.ExchangeFunc {
	responses := map[string]*dns.Msg{}
	set := func(server, qname string, msg *dns.Msg) {
		responses[server+"/"+dns.Fqdn(qname)] = msg
	}
	a := func(name, ip string) dns.RR {
		return &dns.A{
			Hdr: dns.RR_Header{Name: dns.Fqdn(name), Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300},
			A:   net.ParseIP(ip).To4(),
		}
	}
	ns := func(zone, target string) dns.RR {
		return &dns.NS{
			Hdr: dns.RR_Header{Name: dns.Fqdn(zone), Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: 300},
			Ns:  dns.Fqdn(target),
		}
	}

	comRef := new(dns.Msg)
	comRef.Ns = []dns.RR{ns("com", "a.gtld-servers.net")}
	comRef.Extra = []dns.RR{a("a.gtld-servers.net", "192.5.6.30")}
	set("198.41.0.4", "www.example.com", comRef)

	glueless := new(dns.Msg)
	glueless.Ns = []dns.RR{ns("example.com", "ns1.example.net")}
	set("192.5.6.30", "www.example.com", glueless)

	netRef := new(dns.Msg)
	netRef.Ns = []dns.RR{ns("net", "b.gtld-servers.net")}
	netRef.Extra = []dns.RR{a("b.gtld-servers.net", "192.33.14.31")}
	set("198.41.0.4", "ns1.example.net", netRef)

	nsAnswer := new(dns.Msg)
	nsAnswer.Answer = []dns.RR{a("ns1.example.net", "3.3.3.3")}
	set("192.33.14.31", "ns1.example.net", nsAnswer)

	answer := new(dns.Msg)
	answer.Answer = []dns.RR{a("www.example.com", "9.9.9.9")}
	set("3.3.3.3", "www.example.com", answer)

	return func(_ context.Context, server string, msg *dns.Msg, _ bool) (*dns.Msg, error) {
		host := server
		if h, _, err := net.SplitHostPort(server); err == nil {
			host = h
		}
		resp, ok := responses[host+"/"+msg.Question[0].Name]
		if !ok {
			return nil, &net.DNSError{Err: "no mock", Name: msg.Question[0].Name}
		}
		out := resp.Copy()
		out.SetReply(msg)
		out.Answer, out.Ns, out.Extra = resp.Answer, resp.Ns, resp.Extra
		return out, nil
	}
}

func runGlueless(t *testing.T, cfg *Config) string {
	t.Helper()
	tr := traverse.NewTraverser(&traverse.TraverserConfig{
		MaxDepth:  traverse.DefaultMaxDepth,
		QueryType: dns.TypeA,
		RootAddrs: []net.IP{net.ParseIP("198.41.0.4")},
		QueryConfig: &idns.QueryConfig{
			Retries:    1,
			Timeout:    time.Second,
			RetryDelay: time.Millisecond,
		},
	})
	tr.SetExchange(mockGluelessDelegation())

	var buf bytes.Buffer
	formatter := NewFormatter(cfg, &buf)
	if _, err := RunTraversal(context.Background(), tr, cfg, formatter, "www.example.com"); err != nil {
		t.Fatalf("RunTraversal: %v", err)
	}
	return buf.String()
}

func TestResolveProgressHiddenByDefault(t *testing.T) {
	cfg := DefaultConfig()
	cfg.Domain = "www.example.com"
	cfg.QueryType = "a"
	cfg.ShowVersions = false
	cfg.Color = false
	out := runGlueless(t, cfg)

	for _, want := range []string{
		"1.1.1 ns1.example.net -- resolving",
		"1.1.1 ns1.example.net (3.3.3.3)",
		"100.0%: Answer from ns1.example.net (3.3.3.3)",
	} {
		if !strings.Contains(out, want) {
			t.Errorf("output missing %q\n---\n%s", want, out)
		}
	}
	// Resolve subtree nodes (".0." refids) render only under --show-resolves,
	// and resolve outcomes never appear as separate Results entries.
	for _, line := range strings.Split(out, "\n") {
		if strings.HasPrefix(line, "1.1.1.0") {
			t.Errorf("resolve subtree must be hidden by default: %q", line)
		}
	}
	if strings.Contains(out, "ns1.example.net./IN/A") {
		t.Errorf("resolve leaves must not pollute Results\n---\n%s", out)
	}
}

func TestResolveProgressShownWithShowResolves(t *testing.T) {
	cfg := DefaultConfig()
	cfg.Domain = "www.example.com"
	cfg.QueryType = "a"
	cfg.ShowVersions = false
	cfg.ShowResolves = true
	cfg.Color = false
	out := runGlueless(t, cfg)

	for _, want := range []string{
		"1.1.1 ns1.example.net -- resolving",
		"1.1.1.0.1 198.41.0.4 (198.41.0.4)",
		"1.1.1.0.1.1 b.gtld-servers.net (192.33.14.31)",
		"1.1.1 ns1.example.net (3.3.3.3)",
	} {
		if !strings.Contains(out, want) {
			t.Errorf("output missing %q\n---\n%s", want, out)
		}
	}
}

func TestRunTraversalRequiresTraverser(t *testing.T) {
	if _, err := RunTraversal(context.Background(), nil, DefaultConfig(), nil, "example.com"); err == nil {
		t.Fatal("expected error for nil traverser")
	}
}

func TestFormatProbability(t *testing.T) {
	tests := []struct {
		prob float64
		want string
	}{
		{1.0, " 100%"},
		{0.5, "  50%"},
		{0.933, "93.3%"},
		{0.067, " 6.7%"},
		{1.0 / 3, "33.3%"},
	}
	for _, tt := range tests {
		if got := formatProbability(tt.prob); got != tt.want {
			t.Errorf("formatProbability(%v) = %q, want %q", tt.prob, got, tt.want)
		}
	}
}

func TestSummaryStatusLabels(t *testing.T) {
	tests := map[traverse.Status]string{
		traverse.StatusNoData:       "found no such record",
		traverse.StatusReferralLame: "resulted in a lame referral",
		traverse.StatusException:    "resulted in an exception",
		traverse.StatusError:        "resulted in an error",
		traverse.StatusNoGlue:       "found no glue",
		traverse.StatusLoop:         "resulted in a loop",
		traverse.StatusCNAMELoop:    "resulted in a CNAME loop",
		traverse.Status("odd"):      "odd",
	}
	for status, want := range tests {
		if got := summaryStatusLabel(status); got != want {
			t.Errorf("summaryStatusLabel(%q) = %q, want %q", status, got, want)
		}
	}
}

func TestCollectUniqueServerIPs(t *testing.T) {
	servers := map[string][]string{
		"ns1.example.com": {"1.1.1.1", "2.2.2.2"},
		"ns2.example.com": {"1.1.1.1"},
	}
	ips := collectUniqueServerIPs(servers)
	if len(ips) != 2 {
		t.Errorf("unique IPs = %v", ips)
	}
}

func TestReverseString(t *testing.T) {
	if got := reverseString("abc"); got != "cba" {
		t.Errorf("reverseString = %q", got)
	}
}