ExploreDNS/internal/output/formatter_test.go

package output

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

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

func TestDefaultConfig(t *testing.T) {
	cfg := DefaultConfig()
	if !cfg.ShowProgress {
		t.Fatal("expected ShowProgress default true")
	}
	if cfg.Format != FormatText {
		t.Fatalf("Format = %v, want text", cfg.Format)
	}
}

func TestNewFormatterSelectsImplementation(t *testing.T) {
	text := NewFormatter(DefaultConfig(), &bytes.Buffer{})
	if _, ok := text.(*textFormatter); !ok {
		t.Fatalf("expected text formatter, got %T", text)
	}

	jsonCfg := DefaultConfig()
	jsonCfg.Format = FormatJSON
	jsonFmt := NewFormatter(jsonCfg, &bytes.Buffer{})
	if _, ok := jsonFmt.(*jsonFormatter); !ok {
		t.Fatalf("expected json formatter, got %T", jsonFmt)
	}
}

func TestComputeSummaryAggregatesAnswers(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, nil)
	resp := &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("93.184.216.34"),
				},
			},
		},
	}

	stats := ComputeSummary([]traverse.TraversalResult{{Referral: ref, Response: resp}})
	if len(stats.Answers) != 1 {
		t.Fatalf("answers = %d, want 1", len(stats.Answers))
	}
	if stats.Answers[0].Prob != 1.0 {
		t.Fatalf("prob = %v, want 1.0", stats.Answers[0].Prob)
	}
}

func TestTextFormatterSummaryOutput(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, nil)
	resp := &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("93.184.216.34"),
				},
			},
		},
	}

	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Color = false
	cfg.ShowServers = false
	cfg.ShowResults = false
	formatter := NewFormatter(cfg, &buf)

	if err := formatter.WriteSummary([]traverse.TraversalResult{{Referral: ref, Response: resp}}); err != nil {
		t.Fatalf("WriteSummary: %v", err)
	}

	out := buf.String()
	if !strings.Contains(out, "Summary:") {
		t.Fatalf("expected summary header, got %q", out)
	}
	if !strings.Contains(out, "100%") {
		t.Fatalf("expected probability in summary, got %q", out)
	}
	if !strings.Contains(out, "93.184.216.34") {
		t.Fatalf("expected answer IP in summary, got %q", out)
	}
}

func TestJSONFormatterProducesValidOutput(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, nil)
	resp := &traverse.Response{
		Referral: ref,
		Server:   net.ParseIP("198.41.0.4"),
		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("93.184.216.34"),
				},
			},
		},
	}

	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Format = FormatJSON
	cfg.Domain = "example.com"
	cfg.QueryType = "A"
	cfg.ShowServers = false
	cfg.ShowResults = true
	cfg.ShowSummaryResults = true
	formatter := NewFormatter(cfg, &buf)

	if err := formatter.WriteProgress(traverse.TraversalEvent{
		Stage:  traverse.EventStart,
		Result: traverse.TraversalResult{Referral: ref},
	}); err != nil {
		t.Fatalf("WriteProgress: %v", err)
	}
	if err := formatter.WriteSummary([]traverse.TraversalResult{{Referral: ref, Response: resp}}); err != nil {
		t.Fatalf("WriteSummary: %v", err)
	}
	if err := formatter.Flush(); err != nil {
		t.Fatalf("Flush: %v", err)
	}

	var payload map[string]any
	if err := json.Unmarshal(buf.Bytes(), &payload); err != nil {
		t.Fatalf("invalid json: %v\n%s", err, buf.String())
	}
	if payload["domain"] != "example.com" {
		t.Fatalf("domain = %v", payload["domain"])
	}
	if _, ok := payload["summary"]; !ok {
		t.Fatalf("expected summary in json output")
	}
}

func TestTextFormatterWriteSummaryShowsVersions(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, "com.", 0, 1.0, nil)
	serverIP := net.ParseIP("198.41.0.4")
	resp := &traverse.Response{
		Referral: ref,
		Server:   serverIP,
		Type:     traverse.RespAnswer,
	}

	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Color = false
	cfg.ShowVersions = true
	cfg.ShowServers = true
	cfg.ShowResults = false
	cfg.ShowSummaryResults = false
	formatter := NewFormatter(cfg, &buf)
	formatter.SetFingerprints(map[string]string{serverIP.String(): "BIND 9.18.1"})

	if err := formatter.WriteSummary([]traverse.TraversalResult{{Referral: ref, Response: resp}}); err != nil {
		t.Fatalf("WriteSummary: %v", err)
	}

	out := buf.String()
	if !strings.Contains(out, "BIND 9.18.1") {
		t.Fatalf("expected version string in text output, got %q", out)
	}
}

func TestJSONFormatterWriteSummaryShowsVersions(t *testing.T) {
	ref := traverse.NewReferral("example.com.", dns.TypeA, "com.", 0, 1.0, nil)
	serverIP := net.ParseIP("198.41.0.4")
	resp := &traverse.Response{
		Referral: ref,
		Server:   serverIP,
		Type:     traverse.RespAnswer,
	}

	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Format = FormatJSON
	cfg.Domain = "example.com"
	cfg.QueryType = "A"
	cfg.ShowVersions = true
	cfg.ShowServers = true
	cfg.ShowResults = false
	cfg.ShowSummaryResults = false
	formatter := NewFormatter(cfg, &buf)
	formatter.SetFingerprints(map[string]string{serverIP.String(): "Unbound 1.17.0"})

	if err := formatter.WriteSummary([]traverse.TraversalResult{{Referral: ref, Response: resp}}); err != nil {
		t.Fatalf("WriteSummary: %v", err)
	}
	if err := formatter.Flush(); err != nil {
		t.Fatalf("Flush: %v", err)
	}

	var payload map[string]any
	if err := json.Unmarshal(buf.Bytes(), &payload); err != nil {
		t.Fatalf("invalid json: %v\n%s", err, buf.String())
	}
	servers, ok := payload["servers"].([]any)
	if !ok || len(servers) == 0 {
		t.Fatalf("expected servers in json output, got %v", payload)
	}
	srv, ok := servers[0].(map[string]any)
	if !ok {
		t.Fatalf("expected server object, got %T", servers[0])
	}
	if srv["version"] != "Unbound 1.17.0" {
		t.Fatalf("expected version = %q, got %v", "Unbound 1.17.0", srv["version"])
	}
}

func TestRunTraversalUsesHooks(t *testing.T) {
	answerResp := func() *miekgdns.Msg {
		m := new(miekgdns.Msg)
		m.SetReply(new(miekgdns.Msg))
		m.Answer = append(m.Answer, &miekgdns.A{
			Hdr: miekgdns.RR_Header{Name: "example.com.", Rrtype: dns.TypeA, Class: miekgdns.ClassINET, Ttl: 300},
			A:   net.ParseIP("93.184.216.34"),
		})
		return m
	}()

	tr := traverse.NewTraverser(&traverse.TraverserConfig{
		MaxDepth:  5,
		QueryType: dns.TypeA,
		RootAddrs: []net.IP{net.ParseIP("198.41.0.4")},
	})
	tr.SetExchange(func(ctx context.Context, server string, msg *miekgdns.Msg, useTCP bool) (*miekgdns.Msg, error) {
		return answerResp.Copy(), nil
	})

	var buf bytes.Buffer
	cfg := DefaultConfig()
	cfg.Color = false
	cfg.ShowServers = false
	cfg.ShowResults = false
	cfg.ShowSummaryResults = true
	formatter := NewFormatter(cfg, &buf)

	_, err := RunTraversal(context.Background(), tr, cfg, formatter, "example.com")
	if err != nil {
		t.Fatalf("RunTraversal: %v", err)
	}
	if !strings.Contains(buf.String(), "Summary:") {
		t.Fatalf("expected formatted summary output, got %q", buf.String())
	}
}