ExploreDNS/internal/output/formatter_test.go

package output

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

	"gitea.hansenits.com.au/hits/ExploreDNS/internal/dns"
	"gitea.hansenits.com.au/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 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())
	}
}

func TestJSONFormatterWriteResolveAndResult(t *testing.T) {
ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, 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: "example.com.", Rrtype: dns.TypeA, Class: miekgdns.ClassINET},
A:   net.ParseIP("1.2.3.4"),
},
},
},
}

var buf bytes.Buffer
cfg := DefaultConfig()
cfg.Format = FormatJSON
cfg.Domain = "example.com"
cfg.QueryType = "A"
cfg.ShowResolves = true
cfg.ShowAllStats = true
cfg.ShowProgress = true
f := NewFormatter(cfg, &buf).(*jsonFormatter)

// WriteResolve
if err := f.WriteResolve(traverse.TraversalEvent{
Stage:  traverse.EventStart,
Result: traverse.TraversalResult{Referral: ref, Response: resp},
}); err != nil {
t.Fatalf("WriteResolve: %v", err)
}

// WriteResult
if err := f.WriteResult(traverse.TraversalResult{Referral: ref, Response: resp}); err != nil {
t.Fatalf("WriteResult: %v", err)
}

// WriteProgress with EventComplete to cover stageName "complete"
if err := f.WriteProgress(traverse.TraversalEvent{
Stage:  traverse.EventComplete,
Result: traverse.TraversalResult{Referral: ref, Response: resp},
}); err != nil {
t.Fatalf("WriteProgress EventComplete: %v", err)
}

if err := f.Flush(); err != nil {
t.Fatalf("Flush: %v", err)
}
}

func TestJSONFormatterWriteResolveFlagOff(t *testing.T) {
ref := traverse.NewReferral("example.com.", dns.TypeA, ".", 0, 1.0, nil)

var buf bytes.Buffer
cfg := DefaultConfig()
cfg.Format = FormatJSON
cfg.ShowResolves = false
cfg.ShowAllStats = false
f := NewFormatter(cfg, &buf).(*jsonFormatter)

if err := f.WriteResolve(traverse.TraversalEvent{
Stage:  traverse.EventStart,
Result: traverse.TraversalResult{Referral: ref},
}); err != nil {
t.Fatalf("WriteResolve: %v", err)
}
if err := f.WriteResult(traverse.TraversalResult{Referral: ref}); err != nil {
t.Fatalf("WriteResult: %v", err)
}
}

func TestJSONFormatterWriteSummaryWithServers(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,
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: ref, Response: resp}}

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

if err := f.WriteSummary(results); err != nil {
t.Fatalf("WriteSummary: %v", err)
}
if err := f.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 _, ok := payload["servers"]; !ok {
t.Error("expected 'servers' field in JSON output")
}
}

func TestNewFormatterNilWriter(t *testing.T) {
// Should not panic with nil writer
cfg := DefaultConfig()
f := NewFormatter(cfg, nil)
if f == nil {
t.Error("NewFormatter should not return nil")
}
}

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

var buf bytes.Buffer
cfg := DefaultConfig()
cfg.ShowProgress = false
cfg.ShowResolves = true
cfg.ShowAllStats = true
cfg.Color = false
formatter := NewFormatter(cfg, &buf)
hooks := AttachHooks(cfg, formatter)

// Trigger a resolve event
hooks.OnEvent(traverse.TraversalEvent{
Stage:    traverse.EventStart,
IsResolve: true,
Result:   traverse.TraversalResult{Referral: ref, Response: resp},
})

if buf.Len() == 0 {
t.Error("expected resolve output when ShowResolves is true")
}
}

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

var buf bytes.Buffer
cfg := DefaultConfig()
cfg.ShowProgress = false
cfg.ShowResolves = false
cfg.ShowAllStats = true
cfg.Color = false
formatter := NewFormatter(cfg, &buf)
hooks := AttachHooks(cfg, formatter)

hooks.OnEvent(traverse.TraversalEvent{
Stage:  traverse.EventComplete,
Result: traverse.TraversalResult{Referral: ref, Response: resp},
})
}