ExploreDNS/internal/output/formatter.go

package output

import (
	"io"
	"os"

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

type Format int

const (
	FormatText Format = iota
	FormatJSON
)

type Config struct {
	Format             Format
	Domain             string
	QueryType          string
	ShowProgress       bool
	ShowResolves       bool
	ShowServers        bool
	ShowVersions       bool
	ShowAllStats       bool
	ShowResults        bool
	ShowSummaryResults bool
	Verbose            bool
	Quiet              bool
	Color              bool
}

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") == "",
	}
}

type Formatter interface {
	WriteProgress(event traverse.TraversalEvent) error
	WriteResolve(event traverse.TraversalEvent) error
	WriteResult(result traverse.TraversalResult) error
	WriteSummary(results []traverse.TraversalResult) 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
	}
	return &traverse.TraverserHooks{
		OnEvent: func(event traverse.TraversalEvent) {
			switch {
			case event.IsResolve && cfg.ShowResolves:
				_ = formatter.WriteResolve(event)
			case !event.IsResolve && cfg.ShowProgress:
				_ = formatter.WriteProgress(event)
			}
			if event.Stage == traverse.EventComplete && cfg.ShowAllStats {
				_ = formatter.WriteResult(event.Result)
			}
		},
	}
}