ExploreDNS/internal/output/text.go

package output

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

	"github.com/hits/ExploreDNS/internal/dns"
	"github.com/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}
}

func (f *textFormatter) WriteProgress(event traverse.TraversalEvent) error {
	if event.Stage != traverse.EventStart {
		return nil
	}
	line := f.formatReferralLine(event.Result, false)
	if !event.Result.Referral.HasAddresses() {
		line += " -- resolving"
	}
	return f.writeLine(line)
}

func (f *textFormatter) WriteResolve(event traverse.TraversalEvent) error {
	if event.Stage != traverse.EventStart {
		return nil
	}
	return f.writeLine(f.formatReferralLine(event.Result, true))
}

func (f *textFormatter) WriteResult(result traverse.TraversalResult) error {
	if result.Response == nil || result.Referral == nil {
		return nil
	}
	prefix := strings.Repeat("  ", result.Referral.Depth+1)
	line := prefix + f.formatResultLine(result)
	return f.writeLine(line)
}

func (f *textFormatter) WriteSummary(results []traverse.TraversalResult) error {
	if f.cfg.ShowServers {
		if err := f.writeServers(results); err != nil {
			return err
		}
	}

	if f.cfg.ShowResults {
		if err := f.writeResults(results); err != nil {
			return err
		}
	}

	if f.cfg.ShowSummaryResults {
		if err := f.writeSummaryResults(results); err != nil {
			return err
		}
	}

	return nil
}

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

func (f *textFormatter) writeServers(results []traverse.TraversalResult) error {
	servers := collectServers(results)
	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))
	for name := range servers {
		names = append(names, name)
	}
	sort.Slice(names, func(i, j int) bool {
		return strings.ToLower(names[i]) > strings.ToLower(names[j])
	})

	width := 16
	for _, name := range names {
		if len(name) > width {
			width = len(name)
		}
	}

	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, line); err != nil {
				return err
			}
		}
	}
	_, err := fmt.Fprintln(f.w)
	return err
}

func (f *textFormatter) writeResults(results []traverse.TraversalResult) error {
	if _, err := fmt.Fprintln(f.w, "Results:"); err != nil {
		return err
	}

	terminal := terminalResults(results)
	for _, result := range terminal {
		prefix := strings.Repeat("  ", result.Referral.Depth+1)
		line := prefix + f.formatResultLine(result)
		if _, err := fmt.Fprintln(f.w, line); err != nil {
			return err
		}
	}
	_, err := fmt.Fprintln(f.w)
	return err
}

func (f *textFormatter) writeSummaryResults(results []traverse.TraversalResult) error {
	stats := ComputeSummary(results)
	if stats == nil {
		return nil
	}

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

	prefix := "  "
	for _, answer := range stats.Answers {
		line := fmt.Sprintf("%s%s answered with %s", prefix, formatProbability(answer.Prob), answer.RData)
		if _, err := fmt.Fprintln(f.w, f.colorize(line, colorGreen)); err != nil {
			return err
		}
	}

	types := make([]string, 0, len(stats.ByType))
	for respType := range stats.ByType {
		types = append(types, respType)
	}
	sort.Strings(types)

	for _, respType := range types {
		line := fmt.Sprintf("%s%s %s", prefix, formatProbability(stats.ByType[respType]), summaryTypeLabel(respType))
		if _, err := fmt.Fprintln(f.w, line); err != nil {
			return err
		}
	}

	_, err := fmt.Fprintln(f.w)
	return err
}

func (f *textFormatter) formatReferralLine(result traverse.TraversalResult, isResolve bool) string {
	ref := result.Referral
	if ref == nil {
		return ""
	}

	indent := strings.Repeat("  ", ref.Depth)
	refID := referralID(ref)
	server := referralServerLabel(ref, result.Response)
	qtype := dns.QNameType(ref.Qtype)
	qname := trimDomain(ref.Name)

	if f.cfg.Verbose {
		bailiwick := trimDomain(ref.Bailiwick)
		if isResolve {
			return fmt.Sprintf("%s%s [%s] %s <%s>", indent, refID, qname, server, bailiwick)
		}
		return fmt.Sprintf("%s%s [%s] %s <%s> (%s)", indent, refID, qname, server, bailiwick, qtype)
	}

	if isResolve {
		return fmt.Sprintf("%s%s %s", indent, refID, server)
	}
	return fmt.Sprintf("%s%s %s (%s)", indent, refID, server, qtype)
}

func (f *textFormatter) formatResultLine(result traverse.TraversalResult) string {
	prob := formatProbability(result.Referral.Prob)
	switch result.Response.Type {
	case traverse.RespAnswer:
		key, rrs := answerKey(result.Response)
		if key == "" {
			return fmt.Sprintf("%s resulted in answer", prob)
		}
		if len(rrs) == 1 {
			return f.colorize(fmt.Sprintf("%s answered with %s", prob, rrs[0]), colorGreen)
		}
		return f.colorize(fmt.Sprintf("%s answered with %s", prob, strings.Join(rrs, " / ")), colorGreen)
	case traverse.RespNODATA:
		return fmt.Sprintf("%s found no such record", prob)
	case traverse.RespNXDOMAIN:
		return f.colorize(fmt.Sprintf("%s name does not exist", prob), colorYellow)
	case traverse.RespSERVFAIL:
		return f.colorize(fmt.Sprintf("%s resulted in SERVFAIL", prob), colorRed)
	case traverse.RespREFUSED:
		return f.colorize(fmt.Sprintf("%s query refused by server", prob), colorRed)
	case traverse.RespNOTIMPL:
		return f.colorize(fmt.Sprintf("%s query type not implemented by server", prob), colorRed)
	case traverse.RespCNAMELoop:
		msg := "CNAME loop detected"
		if result.Response.ErrorMessage != "" {
			msg = result.Response.ErrorMessage
		}
		return f.colorize(fmt.Sprintf("%s %s", prob, msg), colorRed)
	case traverse.RespError:
		msg := "resulted in an error"
		if result.Response.ErrorMessage != "" {
			msg = result.Response.ErrorMessage
		}
		return f.colorize(fmt.Sprintf("%s %s", prob, msg), colorRed)
	default:
		return fmt.Sprintf("%s %s", prob, result.Response.Type)
	}
}

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
}

func referralID(ref *traverse.Referral) string {
	if ref == nil {
		return ""
	}
	return fmt.Sprintf("%d", ref.Depth+1)
}

func referralServerLabel(ref *traverse.Referral, resp *traverse.Response) string {
	if resp != nil && resp.Server != nil {
		return resp.Server.String()
	}
	if ref.HasAddresses() {
		ips := make([]string, 0, len(ref.Addresses))
		for _, addr := range ref.Addresses {
			ips = append(ips, addr.String())
		}
		return strings.Join(ips, ", ")
	}
	if ref.NSName != "" {
		return ref.NSName
	}
	if ref.Bailiwick != "" && ref.Bailiwick != "." {
		return trimDomain(ref.Bailiwick)
	}
	return "unknown"
}

func terminalResults(results []traverse.TraversalResult) []traverse.TraversalResult {
	var terminal []traverse.TraversalResult
	for _, result := range results {
		if result.Response != nil && result.Response.IsTerminal() {
			terminal = append(terminal, result)
		}
	}
	return terminal
}

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