ExploreDNS/internal/output/text.go

package output

import (
	"fmt"
	"io"
	"sort"
	"strings"

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

type textFormatter struct {
	cfg *Config
	w   io.Writer
}

func newTextFormatter(cfg *Config, w io.Writer) *textFormatter {
	return &textFormatter{cfg: cfg, w: w}
}

// WriteHeader renders the pre-run header block (bin/dnstraverse): the "#"
// settings lines, the initial root, and the "Running query" line. --quiet
// suppresses the whole block. The EDNS0 state reflects the UDP size (the Ruby
// source always printed "on" — a documented deviation we fix).
func (f *textFormatter) WriteHeader(roots []traverse.StartServer) error {
	if f.cfg.Quiet {
		return nil
	}
	var b strings.Builder
	if f.cfg.Fast {
		b.WriteString("# Using fast mode\n")
	}
	if !f.cfg.AllRootServers {
		b.WriteString("# Limiting traverse to one root\n")
	}
	edns := "on"
	if f.cfg.UDPSize <= 512 {
		edns = "off"
	}
	fmt.Fprintf(&b, "# UDP size %d (EDNS0 is %s)\n", f.cfg.UDPSize, edns)
	fmt.Fprintf(&b, "# Retries %d, max depth %d\n", f.cfg.Retries, f.cfg.MaxDepth)
	fmt.Fprintf(&b, "# Allow TCP is %t, always TCP is %t\n", f.cfg.AllowTCP, f.cfg.AlwaysTCP)
	if len(roots) > 0 {
		ip := ""
		if len(roots[0].IPs) > 0 {
			ip = roots[0].IPs[0]
		}
		fmt.Fprintf(&b, "Using %s (%s) as initial root\n", roots[0].Name, ip)
		if f.cfg.AllRootServers {
			b.WriteString("All roots:\n")
			for _, root := range roots {
				fmt.Fprintf(&b, "  %s %s\n", root.Name, strings.Join(root.IPs, ", "))
			}
		}
	}
	fmt.Fprintf(&b, "Running query %s type %s\n", f.cfg.Domain, f.cfg.QueryType)
	_, err := io.WriteString(f.w, b.String())
	return err
}

func (f *textFormatter) WriteProgress(event traverse.TraversalEvent) error {
	ref := event.Referral
	if ref == nil || ref.IsRootRoot() {
		return nil
	}
	switch event.Stage {
	case traverse.StageStart:
		line := f.referralTxt(ref)
		if !ref.Resolved() {
			line += " -- resolving"
		}
		return f.writeLine(line)
	case traverse.StageAnswerFast:
		return f.writeLine(fmt.Sprintf("%s -- completed earlier (%s)",
			f.referralTxt(ref), event.CompletedEarlier))
	case traverse.StageNewReferralSet:
		// One line per extra childset: the parent refid and the IP that
		// produced the children (progress_main :new_referral_set).
		refid := event.RefID
		if i := strings.LastIndex(refid, "."); i >= 0 {
			refid = refid[:i]
		}
		return f.writeLine(fmt.Sprintf("%s %s", refid, ref.ParentIP))
	case traverse.StageAnswer:
		if f.cfg.Verbose {
			for _, warning := range ref.Warnings {
				if err := f.writeLine(fmt.Sprintf("%s WARNING: %s", event.RefID, warning)); err != nil {
					return err
				}
			}
		}
		if f.cfg.ShowAllStats {
			return f.writeStatsBlocks(ref, fmt.Sprintf("%s Results:", event.RefID), false)
		}
	}
	return nil
}

// WriteResolve renders progress for glue-resolution subtree nodes
// (progress_resolves): like main progress, but the fast-mode marker carries
// no refid.
func (f *textFormatter) WriteResolve(event traverse.TraversalEvent) error {
	ref := event.Referral
	if ref == nil || ref.IsRootRoot() {
		return nil
	}
	switch event.Stage {
	case traverse.StageStart:
		return f.writeLine(f.referralTxt(ref))
	case traverse.StageAnswerFast:
		return f.writeLine(f.referralTxt(ref) + " -- completed earlier")
	}
	return nil
}

// referralTxt renders one progress row: "<refid> <server> (<ips>)"; verbose
// adds "[qname]" and "<bailiwick>". Unresolved servers have no parens
// (referral_txt_normal / referral_txt_verbose in bin/dnstraverse).
func (f *textFormatter) referralTxt(ref *traverse.Referral) string {
	var b strings.Builder
	b.WriteString(ref.RefID)
	if f.cfg.Verbose {
		fmt.Fprintf(&b, " [%s]", ref.Qname)
	}
	b.WriteString(" " + ref.Server)
	if ref.Resolved() {
		fmt.Fprintf(&b, " (%s)", ref.TxtIPs())
	}
	if f.cfg.Verbose {
		fmt.Fprintf(&b, " <%s>", ref.Bailiwick)
	}
	return b.String()
}

func (f *textFormatter) WriteSummary(root *traverse.Referral, servers map[string][]string) error {
	// Blank line separating progress from the sections (bin/dnstraverse:
	// "puts if options[:progress]").
	if f.cfg.ShowProgress {
		if _, err := fmt.Fprintln(f.w); err != nil {
			return err
		}
	}
	if f.cfg.ShowServers {
		if err := f.writeServers(servers); err != nil {
			return err
		}
	}
	if f.cfg.ShowResults {
		if err := f.writeResults(root); err != nil {
			return err
		}
	}
	if f.cfg.ShowSummaryResults {
		if err := f.writeSummaryResults(root); err != nil {
			return err
		}
	}
	return nil
}

func (f *textFormatter) Flush() error {
	return nil
}

// writeServers renders "The following servers were encountered:" sorted by
// lowercased reversed name (bin/dnstraverse); the name column is at least 16
// characters wide.
func (f *textFormatter) writeServers(servers map[string][]string) error {
	if len(servers) == 0 {
		return nil
	}

	if _, err := fmt.Fprintln(f.w, "The following servers were encountered:"); err != nil {
		return err
	}

	names := make([]string, 0, len(servers))
	width := 16
	for name := range servers {
		names = append(names, name)
		if len(name) > width {
			width = len(name)
		}
	}
	sort.Slice(names, func(i, j int) bool {
		return reverseString(strings.ToLower(names[i])) < reverseString(strings.ToLower(names[j]))
	})

	for _, name := range names {
		for _, ip := range servers[name] {
			line := fmt.Sprintf("%*s: %-15s", width, name, ip)
			if f.cfg.ShowVersions {
				if version, ok := f.cfg.Fingerprints[ip]; ok && version != "" {
					line += " " + version
				}
			}
			if _, err := fmt.Fprintln(f.w, strings.TrimRight(line, " ")); err != nil {
				return err
			}
		}
	}
	_, err := fmt.Fprintln(f.w)
	return err
}

func reverseString(s string) string {
	runes := []rune(s)
	for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
		runes[i], runes[j] = runes[j], runes[i]
	}
	return string(runes)
}

func (f *textFormatter) writeResults(root *traverse.Referral) error {
	if root == nil {
		return nil
	}
	if _, err := fmt.Fprintln(f.w, "Results:"); err != nil {
		return err
	}
	if err := f.writeStatsBlocks(root, "", true); err != nil {
		return err
	}
	_, err := fmt.Fprintln(f.w)
	return err
}

// writeStatsBlocks renders every aggregated leaf of ref, sorted by stats key
// (referral.rb stats_display). With spacing, blocks are separated by blank
// lines.
func (f *textFormatter) writeStatsBlocks(ref *traverse.Referral, prefix string, spacing bool) error {
	first := true
	for _, leaf := range ref.StatsList() {
		if spacing && !first {
			if _, err := fmt.Fprintln(f.w); err != nil {
				return err
			}
		}
		first = false
		for _, line := range f.formatLeaf(leaf, prefix) {
			if err := f.writeLine(line); err != nil {
				return err
			}
		}
	}
	return nil
}

func (f *textFormatter) writeSummaryResults(root *traverse.Referral) error {
	stats := root.SummaryStats()
	if stats == nil {
		return nil
	}

	if _, err := fmt.Fprintln(f.w, "Summary Results:"); err != nil {
		return err
	}

	prefix := "  "
	for _, answer := range stats.Answers {
		initial := fmt.Sprintf("%s%s answered with ", prefix, formatProbability(answer.Prob))
		var rrs []string
		for _, rr := range answer.RRs {
			rrs = append(rrs, collapseWhitespace(rr.String()))
		}
		line := initial + strings.Join(rrs, "\n"+strings.Repeat(" ", len(initial)))
		if _, err := fmt.Fprintln(f.w, f.colorize(line, colorGreen)); err != nil {
			return err
		}
	}

	statuses := make([]traverse.Status, 0, len(stats.ByStatus))
	for status := range stats.ByStatus {
		if status != traverse.StatusAnswered {
			statuses = append(statuses, status)
		}
	}
	sort.Slice(statuses, func(i, j int) bool { return statuses[i] < statuses[j] })

	for _, status := range statuses {
		line := fmt.Sprintf("%s%s %s", prefix, formatProbability(stats.ByStatus[status]), summaryStatusLabel(status))
		if _, err := fmt.Fprintln(f.w, line); err != nil {
			return err
		}
	}

	return nil
}

// formatLeaf renders one aggregated leaf per referral.rb stats_display:
// "%5.1f%%: <wording>" plus indented RRs for answers and the
// "While querying" line when the failing query differs from the original.
func (f *textFormatter) formatLeaf(leaf *traverse.StatsEntry, prefix string) []string {
	resp := leaf.Response
	ref := leaf.Referral
	if resp == nil || ref == nil {
		return nil
	}
	indent := prefix + strings.Repeat(" ", 12)
	where := fmt.Sprintf("%s (%s)", ref.Server, resp.IP)
	head := fmt.Sprintf("%s%5.1f%%: ", prefix, leaf.Prob*100)

	var lines []string
	switch resp.Status {
	case traverse.StatusException:
		msg := ""
		if resp.DQ != nil {
			msg = resp.DQ.ExceptionMessage
		}
		lines = append(lines, head+f.colorize(fmt.Sprintf("%s at %s", msg, where), colorRed))
	case traverse.StatusNoGlue:
		parent := ""
		if ref.Parent != nil {
			parent = ref.Parent.Server
		}
		lines = append(lines, head+f.colorize(fmt.Sprintf("No glue at %s (%s) for %s", parent, resp.IP, ref.Server), colorYellow))
	case traverse.StatusReferralLame:
		parent := ""
		if ref.Parent != nil {
			parent = ref.Parent.Server
		}
		lines = append(lines, head+f.colorize(fmt.Sprintf("Lame referral from %s (%s) to %s", parent, ref.ParentIP, where), colorYellow))
	case traverse.StatusLoop:
		lines = append(lines, head+f.colorize(fmt.Sprintf("Loop encountered at %s", resp.Server), colorRed))
	case traverse.StatusCNAMELoop:
		lines = append(lines, head+f.colorize(fmt.Sprintf("CNAME loop encountered at %s", resp.Server), colorRed))
	case traverse.StatusError:
		msg := ""
		if resp.DQ != nil {
			msg = resp.DQ.ErrorMessage
		}
		lines = append(lines, head+f.colorize(fmt.Sprintf("%s at %s", msg, where), colorRed))
	case traverse.StatusNoData:
		lines = append(lines, head+fmt.Sprintf("NODATA (for this type) at %s", where))
	case traverse.StatusAnswered:
		lines = append(lines, head+f.colorize(fmt.Sprintf("Answer from %s", where), colorGreen))
		if resp.DQ != nil {
			for _, rr := range resp.DQ.Answers {
				lines = append(lines, indent+rr.String())
			}
		}
	default:
		// The Ruby fallback prints "Stopped at <server> (<ip>))" with a stray
		// paren — a documented deviation we fix.
		lines = append(lines, head+fmt.Sprintf("Stopped at %s", where))
		lines = append(lines, indent+leaf.Key)
	}

	if resp.Status != traverse.StatusAnswered {
		origQname, origQclass, origQtype := originalQuery(ref)
		if resp.Qname != origQname || resp.Qclass != origQclass || resp.Qtype != origQtype {
			lines = append(lines, indent+fmt.Sprintf("While querying %s/%s/%s",
				resp.Qname, traverse.ClassToString(resp.Qclass), traverse.TypeToString(resp.Qtype)))
		}
	}
	return lines
}

// originalQuery walks to the rootroot node to find the query the whole
// traversal was started for.
func originalQuery(ref *traverse.Referral) (string, uint16, uint16) {
	top := ref
	for top.Parent != nil {
		top = top.Parent
	}
	return top.Qname, top.Qclass, top.Qtype
}

func (f *textFormatter) writeLine(line string) error {
	if line == "" {
		return nil
	}
	_, err := fmt.Fprintln(f.w, line)
	return err
}

func (f *textFormatter) colorize(text, color string) string {
	if !f.cfg.Color || color == "" {
		return text
	}
	return color + text + colorReset
}

const (
	colorReset  = "\033[0m"
	colorGreen  = "\033[32m"
	colorYellow = "\033[33m"
	colorRed    = "\033[31m"
)