ExploreDNS/internal/output/formatter.go

// Package output renders ExploreDNS traversal results for human consumption
// or machine processing.
//
// Two formats are supported:
//
//   - FormatText — dnstraverse-style header/progress/results/summary text
//     (default)
//   - FormatJSON — a single JSON document with the aggregated results
//
// Create a Formatter via NewFormatter and call RunTraversal to drive the
// traversal engine and stream output incrementally.
package output

import (
	"fmt"
	"io"
	"os"

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

type Format int

const (
	FormatText Format = iota
	FormatJSON
)

type Config struct {
	Format    Format
	Domain    string
	QueryType string

	// Engine settings echoed in the header block (bin/dnstraverse).
	Fast           bool
	AllRootServers bool
	UDPSize        int
	Retries        int
	MaxDepth       int
	AllowTCP       bool
	AlwaysTCP      bool

	ShowProgress       bool
	ShowResolves       bool
	ShowServers        bool
	ShowVersions       bool
	ShowAllStats       bool
	ShowResults        bool
	ShowSummaryResults bool
	Verbose            bool
	Quiet              bool
	Color              bool
	Debug              int

	// Fingerprints maps server IP strings to their version.bind version strings.
	// Populated by RunTraversal when ShowVersions and ShowServers are both true.
	Fingerprints map[string]string
}

func DefaultConfig() *Config {
	return &Config{
		Format:             FormatText,
		Fast:               true,
		UDPSize:            2048,
		Retries:            2,
		MaxDepth:           20,
		AllowTCP:           true,
		ShowProgress:       true,
		ShowResolves:       false,
		ShowServers:        false,
		ShowVersions:       true,
		ShowAllStats:       false,
		ShowResults:        true,
		ShowSummaryResults: true,
		Color:              ColorEnabled(os.Stdout),
	}
}

// ColorEnabled reports whether colour output should be used for w: only when
// NO_COLOR is unset and w is a terminal.
func ColorEnabled(w io.Writer) bool {
	if os.Getenv("NO_COLOR") != "" {
		return false
	}
	f, ok := w.(*os.File)
	if !ok {
		return false
	}
	info, err := f.Stat()
	if err != nil {
		return false
	}
	return info.Mode()&os.ModeCharDevice != 0
}

// Formatter renders traversal progress and the aggregated results. The header
// is written once before the run, progress arrives via the traverser hooks,
// and the aggregated leaves and the servers seen arrive once the run
// completed.
type Formatter interface {
	// WriteHeader renders the pre-run header block from the discovered roots
	// (suppressed entirely by --quiet in text mode).
	WriteHeader(roots []traverse.StartServer) error
	WriteProgress(event traverse.TraversalEvent) error
	WriteResolve(event traverse.TraversalEvent) error
	WriteSummary(root *traverse.Referral, servers map[string][]string) error
	Flush() error
}

func NewFormatter(cfg *Config, w io.Writer) Formatter {
	if cfg == nil {
		cfg = DefaultConfig()
	}
	if w == nil {
		w = os.Stdout
	}
	if cfg.Format == FormatJSON {
		return newJSONFormatter(cfg, w)
	}
	return newTextFormatter(cfg, w)
}

func AttachHooks(cfg *Config, formatter Formatter) *traverse.TraverserHooks {
	if cfg == nil || formatter == nil {
		return nil
	}
	if !cfg.ShowProgress {
		// Ruby registers no progress callbacks at all without show-progress;
		// resolve display additionally requires show-resolves.
		return nil
	}
	logErr := func(context string, err error) {
		if err != nil && cfg.Debug > 0 {
			fmt.Fprintf(os.Stderr, "Debug: formatter %s: %v\n", context, err)
		}
	}
	return &traverse.TraverserHooks{
		OnEvent: func(event traverse.TraversalEvent) {
			if event.IsResolve {
				if cfg.ShowResolves {
					logErr("WriteResolve", formatter.WriteResolve(event))
				}
				return
			}
			logErr("WriteProgress", formatter.WriteProgress(event))
		},
	}
}