ExploreDNS/internal/output/stats.go

package output

import (
	"fmt"
	"sort"
	"strings"

	"github.com/hits/ExploreDNS/internal/dns"
	"github.com/hits/ExploreDNS/internal/traverse"
	miekgdns "github.com/miekg/dns"
)

type summaryEntry struct {
	Type string
	Prob float64
}

type answerEntry struct {
	RData string
	Prob  float64
	RRs   []string
}

// SummaryStats holds the aggregated probability statistics computed from a set
// of traversal results.
type SummaryStats struct {
	// ByType maps ResponseType strings (e.g. "nxdomain", "servfail") to their
	// cumulative probability weight.
	ByType map[string]float64
	// Answers contains per-RDATA probability statistics for answer results.
	Answers []answerEntry
}

// ComputeSummary aggregates terminal TraversalResults into a SummaryStats.
// Returns nil when there are no terminal results to summarise.
func ComputeSummary(results []traverse.TraversalResult) *SummaryStats {
	stats := &SummaryStats{
		ByType: make(map[string]float64),
	}

	for _, result := range results {
		if result.Response == nil || !result.Response.IsTerminal() {
			continue
		}
		if result.Referral == nil {
			continue
		}

		prob := result.Referral.Prob
		respType := result.Response.Type.String()

		switch result.Response.Type {
		case traverse.RespAnswer:
			key, rrs := answerKey(result.Response)
			if key == "" {
				stats.ByType[respType] += prob
				continue
			}
			found := false
			for i := range stats.Answers {
				if stats.Answers[i].RData == key {
					stats.Answers[i].Prob += prob
					found = true
					break
				}
			}
			if !found {
				stats.Answers = append(stats.Answers, answerEntry{
					RData: key,
					Prob:  prob,
					RRs:   rrs,
				})
			}
		default:
			stats.ByType[respType] += prob
		}
	}

	sort.Slice(stats.Answers, func(i, j int) bool {
		return stats.Answers[i].RData < stats.Answers[j].RData
	})

	if len(stats.Answers) == 0 && len(stats.ByType) == 0 {
		return nil
	}
	return stats
}

func answerKey(resp *traverse.Response) (string, []string) {
	if resp == nil || resp.Decoded == nil {
		return "", nil
	}

	var rdatas []string
	var formatted []string
	for _, rr := range resp.Decoded.Answers {
		if _, ok := rr.(*miekgdns.CNAME); ok {
			continue
		}
		rdata := rrDataString(rr)
		if rdata == "" {
			continue
		}
		rdatas = append(rdatas, rdata)
		formatted = append(formatted, dns.FormatRecord(rr))
	}

	if len(rdatas) == 0 {
		return "", nil
	}
	sort.Strings(rdatas)
	return strings.Join(rdatas, " / "), formatted
}

func rrDataString(rr miekgdns.RR) string {
	switch v := rr.(type) {
	case *miekgdns.A:
		return v.A.String()
	case *miekgdns.AAAA:
		return v.AAAA.String()
	case *miekgdns.CNAME:
		return v.Target
	case *miekgdns.NS:
		return v.Ns
	case *miekgdns.MX:
		return fmt.Sprintf("%d %s", v.Preference, v.Mx)
	case *miekgdns.TXT:
		return strings.Join(v.Txt, " ")
	default:
		return rr.String()
	}
}

func formatProbability(prob float64) string {
	text := fmt.Sprintf("%.1f%%", prob*100)
	text = strings.Replace(text, ".0%", "%", 1)
	return fmt.Sprintf("%5s", text)
}

func summaryTypeLabel(respType string) string {
	switch respType {
	case "nodata":
		return "found no such record"
	case "nxdomain":
		return "name does not exist"
	case "servfail":
		return "resulted in SERVFAIL"
	case "refused":
		return "query refused by server"
	case "notimp":
		return "query type not implemented by server"
	case "cname_loop":
		return "resulted in a CNAME loop"
	case "error":
		return "resulted in an error"
	case "referral":
		return "resulted in a referral"
	default:
		return respType
	}
}

func trimDomain(name string) string {
	return strings.TrimSuffix(name, ".")
}

func collectServers(results []traverse.TraversalResult) map[string][]string {
	servers := make(map[string][]string)
	for _, result := range results {
		if result.Response == nil || result.Response.Server == nil {
			continue
		}
		name := serverName(result)
		ip := result.Response.Server.String()
		if containsString(servers[name], ip) {
			continue
		}
		servers[name] = append(servers[name], ip)
	}
	return servers
}

func serverName(result traverse.TraversalResult) string {
	if result.Referral != nil && result.Referral.Bailiwick != "" && result.Referral.Bailiwick != "." {
		return trimDomain(result.Referral.Bailiwick)
	}
	if result.Referral != nil && result.Referral.NSName != "" {
		return result.Referral.NSName
	}
	if result.Response != nil && result.Response.Server != nil {
		return result.Response.Server.String()
	}
	return "unknown"
}

func containsString(items []string, target string) bool {
	for _, item := range items {
		if item == target {
			return true
		}
	}
	return false
}