ExploreDNS/internal/output/stats_test.go

package output

import (
	"net"
	"testing"

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

func makeAnswerResult(name string, ip string, prob float64) traverse.TraversalResult {
	ref := traverse.NewReferral(name, dns.TypeA, ".", 0, prob, nil)
	server := net.ParseIP("198.41.0.4")
	resp := &traverse.Response{
		Referral: ref,
		Server:   server,
		Type:     traverse.RespAnswer,
		Decoded: &dns.DecodedResponse{
			Answers: []miekgdns.RR{
				&miekgdns.A{
					Hdr: miekgdns.RR_Header{Name: name + ".", Rrtype: dns.TypeA, Class: miekgdns.ClassINET},
					A:   net.ParseIP(ip),
				},
			},
		},
	}
	return traverse.TraversalResult{Referral: ref, Response: resp}
}

func TestRRDataStringAllTypes(t *testing.T) {
	tests := []struct {
		rr   miekgdns.RR
		want string
	}{
		{
			&miekgdns.A{Hdr: miekgdns.RR_Header{Rrtype: dns.TypeA}, A: net.ParseIP("1.2.3.4")},
			"1.2.3.4",
		},
		{
			&miekgdns.AAAA{Hdr: miekgdns.RR_Header{Rrtype: miekgdns.TypeAAAA}, AAAA: net.ParseIP("::1")},
			"::1",
		},
		{
			&miekgdns.CNAME{Hdr: miekgdns.RR_Header{Rrtype: miekgdns.TypeCNAME}, Target: "example.com."},
			"example.com.",
		},
		{
			&miekgdns.NS{Hdr: miekgdns.RR_Header{Rrtype: miekgdns.TypeNS}, Ns: "ns1.example.com."},
			"ns1.example.com.",
		},
		{
			&miekgdns.MX{Hdr: miekgdns.RR_Header{Rrtype: miekgdns.TypeMX}, Preference: 10, Mx: "mail.example.com."},
			"10 mail.example.com.",
		},
		{
			&miekgdns.TXT{Hdr: miekgdns.RR_Header{Rrtype: miekgdns.TypeTXT}, Txt: []string{"v=spf1", "include:example.com"}},
			"v=spf1 include:example.com",
		},
	}

	for _, tc := range tests {
		got := rrDataString(tc.rr)
		if got != tc.want {
			t.Errorf("rrDataString(%T) = %q, want %q", tc.rr, got, tc.want)
		}
	}
}

func TestRRDataStringDefault(t *testing.T) {
	// SOA record hits the default case
	rr := &miekgdns.SOA{
		Hdr:  miekgdns.RR_Header{Name: ".", Rrtype: miekgdns.TypeSOA, Class: miekgdns.ClassINET},
		Ns:   "a.root-servers.net.",
		Mbox: "nstld.verisign-grs.com.",
	}
	got := rrDataString(rr)
	if got == "" {
		t.Error("rrDataString(SOA) should return non-empty string via default case")
	}
}

func TestSummaryTypeLabelAllTypes(t *testing.T) {
	cases := map[string]string{
		"nodata":     "found no such record",
		"nxdomain":   "name does not exist",
		"servfail":   "resulted in SERVFAIL",
		"refused":    "query refused by server",
		"notimp":     "query type not implemented by server",
		"cname_loop": "resulted in a CNAME loop",
		"ns_error":   "nameserver lookup failed",
		"error":      "resulted in an error",
		"referral":   "resulted in a referral",
		"unknown_type": "unknown_type",
	}
	for input, want := range cases {
		got := summaryTypeLabel(input)
		if got != want {
			t.Errorf("summaryTypeLabel(%q) = %q, want %q", input, got, want)
		}
	}
}

func TestCollectServersEmpty(t *testing.T) {
	servers := collectServers(nil)
	if len(servers) != 0 {
		t.Errorf("collectServers(nil) = %v, want empty", servers)
	}
}

func TestCollectServersDeduplication(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, "com.", 1, 1.0, nil)
	server := net.ParseIP("1.2.3.4")
	resp := &traverse.Response{
		Referral: ref,
		Server:   server,
		Type:     traverse.RespAnswer,
	}
	result := traverse.TraversalResult{Referral: ref, Response: resp}

	servers := collectServers([]traverse.TraversalResult{result, result})
	name := "com"
	ips := servers[name]
	if len(ips) != 1 {
		t.Errorf("expected deduplication: got %d IPs, want 1", len(ips))
	}
}

func TestCollectServersWithBailiwick(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, "com.", 1, 1.0, nil)
	server := net.ParseIP("1.2.3.4")
	resp := &traverse.Response{
		Referral: ref,
		Server:   server,
		Type:     traverse.RespAnswer,
	}
	result := traverse.TraversalResult{Referral: ref, Response: resp}

	servers := collectServers([]traverse.TraversalResult{result})
	if len(servers) == 0 {
		t.Fatal("expected at least one server entry")
	}
	if _, ok := servers["com"]; !ok {
		t.Errorf("expected server name 'com', got keys: %v", servers)
	}
}

func TestServerNameFallbacks(t *testing.T) {
	// No bailiwick, no NSName, with server IP
	ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, nil)
	resp := &traverse.Response{
		Referral: ref,
		Server:   net.ParseIP("1.2.3.4"),
		Type:     traverse.RespAnswer,
	}
	result := traverse.TraversalResult{Referral: ref, Response: resp}
	name := serverName(result)
	if name != "1.2.3.4" {
		t.Errorf("serverName with root bailiwick = %q, want IP", name)
	}
}

func TestServerNameWithNSName(t *testing.T) {
	ref := &traverse.Referral{
		Name:      "example.com.",
		Qtype:     dns.TypeA,
		Bailiwick: ".",
		NSName:    "ns1.example.com.",
	}
	resp := &traverse.Response{
		Referral: ref,
		Server:   net.ParseIP("5.5.5.5"),
		Type:     traverse.RespAnswer,
	}
	result := traverse.TraversalResult{Referral: ref, Response: resp}
	// Bailiwick is "." so falls through to NSName
	name := serverName(result)
	if name == "" {
		t.Error("serverName should return non-empty string")
	}
}

func TestServerNameNilReferral(t *testing.T) {
	resp := &traverse.Response{
		Server: net.ParseIP("1.2.3.4"),
		Type:   traverse.RespAnswer,
	}
	result := traverse.TraversalResult{Referral: nil, Response: resp}
	name := serverName(result)
	if name == "" {
		t.Error("serverName with nil referral should return non-empty string")
	}
}

func TestContainsString(t *testing.T) {
	items := []string{"a", "b", "c"}
	if !containsString(items, "b") {
		t.Error("containsString should find 'b' in slice")
	}
	if containsString(items, "d") {
		t.Error("containsString should not find 'd' in slice")
	}
	if containsString(nil, "a") {
		t.Error("containsString on nil slice should return false")
	}
}

func TestComputeSummaryMixedResults(t *testing.T) {
	refAnswer := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 0.6, nil)
	respAnswer := &traverse.Response{
		Referral: refAnswer,
		Type:     traverse.RespAnswer,
		Decoded: &dns.DecodedResponse{
			Answers: []miekgdns.RR{
				&miekgdns.A{
					Hdr: miekgdns.RR_Header{Name: "example.com.", Rrtype: dns.TypeA, Class: miekgdns.ClassINET},
					A:   net.ParseIP("1.2.3.4"),
				},
			},
		},
	}

	refNXD := traverse.NewReferral("notexist.com.", dns.TypeA, ".", 0, 0.4, nil)
	respNXD := &traverse.Response{
		Referral: refNXD,
		Type:     traverse.RespNXDOMAIN,
	}

	results := []traverse.TraversalResult{
		{Referral: refAnswer, Response: respAnswer},
		{Referral: refNXD, Response: respNXD},
	}

	stats := ComputeSummary(results)
	if stats == nil {
		t.Fatal("ComputeSummary returned nil for non-empty results")
	}
	if len(stats.Answers) != 1 {
		t.Errorf("expected 1 answer entry, got %d", len(stats.Answers))
	}
	if _, ok := stats.ByType["nxdomain"]; !ok {
		t.Error("expected nxdomain in ByType")
	}
}

func TestComputeSummaryAnswerWithCNAMEOnly(t *testing.T) {
	// Answer with only CNAME record - no final answer, should be in ByType
	ref := traverse.NewReferral("www.example.com.", dns.TypeA, ".", 0, 1.0, nil)
	resp := &traverse.Response{
		Referral: ref,
		Type:     traverse.RespAnswer,
		Decoded: &dns.DecodedResponse{
			Answers: []miekgdns.RR{
				&miekgdns.CNAME{
					Hdr:    miekgdns.RR_Header{Name: "www.example.com.", Rrtype: miekgdns.TypeCNAME, Class: miekgdns.ClassINET},
					Target: "example.com.",
				},
			},
		},
	}
	results := []traverse.TraversalResult{{Referral: ref, Response: resp}}
	stats := ComputeSummary(results)
	if stats == nil {
		t.Fatal("ComputeSummary returned nil")
	}
}

func TestComputeSummaryAccumulates(t *testing.T) {
	// Two answers with the same IP should accumulate probability
	ref1 := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 0.5, nil)
	ref2 := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 0.5, nil)

	makeResp := func(ref *traverse.Referral) *traverse.Response {
		return &traverse.Response{
			Referral: ref,
			Type:     traverse.RespAnswer,
			Decoded: &dns.DecodedResponse{
				Answers: []miekgdns.RR{
					&miekgdns.A{
						Hdr: miekgdns.RR_Header{Name: "example.com.", Rrtype: dns.TypeA, Class: miekgdns.ClassINET},
						A:   net.ParseIP("1.2.3.4"),
					},
				},
			},
		}
	}

	results := []traverse.TraversalResult{
		{Referral: ref1, Response: makeResp(ref1)},
		{Referral: ref2, Response: makeResp(ref2)},
	}
	stats := ComputeSummary(results)
	if stats == nil {
		t.Fatal("ComputeSummary returned nil")
	}
	if len(stats.Answers) != 1 {
		t.Fatalf("expected 1 answer after accumulation, got %d", len(stats.Answers))
	}
	if stats.Answers[0].Prob < 0.99 {
		t.Errorf("accumulated prob = %.2f, want ~1.0", stats.Answers[0].Prob)
	}
}

func TestCollectUniqueServerIPs(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, nil)
	ip1 := net.ParseIP("1.2.3.4")
	ip2 := net.ParseIP("5.6.7.8")

	results := []traverse.TraversalResult{
		{Referral: ref, Response: &traverse.Response{Server: ip1, Type: traverse.RespAnswer}},
		{Referral: ref, Response: &traverse.Response{Server: ip1, Type: traverse.RespAnswer}}, // dup
		{Referral: ref, Response: &traverse.Response{Server: ip2, Type: traverse.RespAnswer}},
		{Referral: ref, Response: nil}, // nil response
	}

	ips := collectUniqueServerIPs(results)
	if len(ips) != 2 {
		t.Errorf("expected 2 unique IPs, got %d", len(ips))
	}
}

func TestCollectUniqueServerIPsEmpty(t *testing.T) {
	ips := collectUniqueServerIPs(nil)
	if len(ips) != 0 {
		t.Errorf("expected 0 IPs for nil results, got %d", len(ips))
	}
}