ExploreDNS/internal/output/formatter.go

package output

import (
	"fmt"
	"io"
	"os"

	"github.com/hits/ExploreDNS/internal/traverse"
)

// Format selects the output format.
type Format int

// Output format constants.
const (
	FormatText Format = iota // human-readable coloured text (default)
	FormatJSON               // machine-readable JSON
)

// Config carries settings that control what the formatter emits.
type Config struct {
	// Format selects text or JSON output.
	Format Format
	// Domain is the queried domain name, included in JSON output.
	Domain string
	// QueryType is the record type queried, included in JSON output.
	QueryType string
	// ShowProgress enables per-referral progress lines.
	ShowProgress bool
	// ShowResolves enables nameserver resolution detail lines.
	ShowResolves bool
	// ShowServers enables the server list section.
	ShowServers bool
	// ShowVersions enables version.bind fingerprint display alongside servers.
	ShowVersions bool
	// ShowAllStats enables per-result stats as they arrive (in text mode).
	ShowAllStats bool
	// ShowResults enables the terminal-results section.
	ShowResults bool
	// ShowSummaryResults enables the probability summary section.
	ShowSummaryResults bool
	// Verbose enables additional detail in progress lines.
	Verbose bool
	// Quiet suppresses the introductory banner line.
	Quiet bool
	// Color enables ANSI terminal colour codes in text output.
	Color bool
	// Debug controls debug verbosity for the formatter itself.
	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
}

// DefaultConfig returns a Config with all output sections enabled, text format,
// and color determined by the NO_COLOR environment variable.
func DefaultConfig() *Config {
	return &Config{
		Format:             FormatText,
		ShowProgress:       true,
		ShowResolves:       true,
		ShowServers:        true,
		ShowVersions:       true,
		ShowAllStats:       true,
		ShowResults:        true,
		ShowSummaryResults: true,
		Color:              os.Getenv("NO_COLOR") == "",
	}
}

// Formatter is the interface that both the text and JSON output backends implement.
// Each method is called by the traversal hooks or by RunTraversal.
type Formatter interface {
	// WriteProgress is called at EventStart for each Referral.
	WriteProgress(event traverse.TraversalEvent) error
	// WriteResolve is called at EventStart for nameserver resolution sub-steps.
	WriteResolve(event traverse.TraversalEvent) error
	// WriteResult is called at EventComplete for each TraversalResult.
	WriteResult(result traverse.TraversalResult) error
	// WriteSummary is called once after all results are collected.
	WriteSummary(results []traverse.TraversalResult) error
	// Flush finalises output (e.g. writes buffered JSON to the writer).
	Flush() error
}

// NewFormatter returns the appropriate Formatter (text or JSON) based on cfg.Format.
// A nil cfg uses DefaultConfig.  A nil w uses os.Stdout.
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)
}

// AttachHooks creates a TraverserHooks that routes traversal events to formatter
// according to cfg visibility settings.  Returns nil when cfg or formatter is nil.
func AttachHooks(cfg *Config, formatter Formatter) *traverse.TraverserHooks {
	if cfg == nil || formatter == nil {
		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) {
			switch {
			case event.IsResolve && cfg.ShowResolves:
				logErr("WriteResolve", formatter.WriteResolve(event))
			case !event.IsResolve && cfg.ShowProgress:
				logErr("WriteProgress", formatter.WriteProgress(event))
			}
			if event.Stage == traverse.EventComplete && cfg.ShowAllStats {
				logErr("WriteResult", formatter.WriteResult(event.Result))
			}
		},
	}
}