Improved Results Display with Grouping and Deduplication

cmd/hednsextractor/hednsextractor.go

@@ -4,118 +4,165 @@ import (
 	"fmt"
 	"regexp"
 	"strconv"
+	"strings"
 
 	"github.com/HuntDownProject/hednsextractor/utils"
+	"github.com/fatih/color"
 	"github.com/projectdiscovery/gologger"
 )
 
 var (
-	outputs []string
+	outputs    = make(map[string][]string) // Map to group domains by IP
+	ptrTracker = make(map[string]bool)     // Map to track PTRs that have already been added
 )
 
 func main() {
-
-	// Parse the stdin
 	utils.ParseStdin()
-
-	// Load parameters from command line and configuration file
 	utils.LoadParameters()
 
-	// Show Banner
-	utils.ShowBanner()
+	if !utils.OptionCmd.Silent {
+		utils.ShowBanner()
+	}
 
-	// read the Workflow from yaml
 	var workflow utils.Workflow
 	if utils.OptionCmd.Workflow != "" {
 		workflow.GetConf(utils.OptionCmd.Workflow)
 
-		for i := range workflow.Domains {
-			utils.IdentifyTarget(workflow.Domains[i])
-		}
-
-		for i := range workflow.Ipaddrs {
-			utils.IdentifyTarget(workflow.Ipaddrs[i])
-		}
-
-		for i := range workflow.Networks {
-			utils.IdentifyTarget(workflow.Networks[i])
-		}
+		// Simplify workflow target processing
+		processTargets(workflow.Domains)
+		processTargets(workflow.Ipaddrs)
+		processTargets(workflow.Networks)
 	}
 
 	hurricane := utils.Hurricane{}
 	hurricane.RunCrawler()
 
 	if utils.OptionCmd.Vtscore && !utils.OptionCmd.Silent {
-		gologger.Info().Msgf("Filtering with Virustotal with a mininum score %s", utils.OptionCmd.VtscoreValue)
+		gologger.Info().Msgf("Filtering with Virustotal, minimum score: %s", utils.OptionCmd.VtscoreValue)
 	}
 
+	// Process results
 	for _, result := range utils.Results {
-		var bMatchedPTR = false
-		var bMatchedDomain = false
+		processResult(result, workflow)
+	}
 
-		if workflow.Regex != "" {
-			var re = regexp.MustCompile(workflow.Regex)
-			bMatchedDomain = re.MatchString(result.Domain)
-			bMatchedPTR = re.MatchString(result.PTR)
-		} else {
-			bMatchedPTR = true
-			bMatchedDomain = true
-		}
+	displayOutputs() // Display results grouped by IP
+}
 
-		if !bMatchedDomain && !bMatchedPTR {
-			continue
-		}
+// Processes workflow targets such as domains, IP addresses, or networks
+func processTargets(targets []string) {
+	for _, target := range targets {
+		utils.IdentifyTarget(target)
+	}
+}
 
-		if utils.OptionCmd.Vtscore {
-			virustotal := utils.Virustotal{}
-			result.VtScore = virustotal.GetVtReport(result.Domain)
-			if score, err := strconv.ParseUint(utils.OptionCmd.VtscoreValue, 10, 64); err == nil {
-				if result.VtScore < score {
-					continue
-				}
-			} else {
-				gologger.Fatal().Msg("Invalid parameter value for vt-score")
-			}
-		}
+// Processes each result and applies regex and VirusTotal score filters
+func processResult(result utils.Result, workflow utils.Workflow) {
+	bMatchedDomain, bMatchedPTR := matchResultWithRegex(result, workflow.Regex)
+
+	// Skip if no domain or PTR matches the regex
+	if !bMatchedDomain && !bMatchedPTR {
+		return
+	}
+
+	// Apply VirusTotal filter if enabled
+	if utils.OptionCmd.Vtscore && !filterByVtScore(result) {
+		return
+	}
 
-		var output = prepareOutput(result, bMatchedDomain, bMatchedPTR)
-		if !utils.Contains(outputs, output) {
-			outputs = append(outputs, output)
+	// Group domains by IP
+	if bMatchedDomain && result.Domain != "" {
+		formatAndStoreResult(result.IPAddr, result.Domain, "Domain")
+	}
+
+	// Avoid PTR duplicates
+	if bMatchedPTR && result.PTR != "" {
+		ptrKey := fmt.Sprintf("%s:%s", result.IPAddr, result.PTR)
+		if !ptrTracker[ptrKey] {
+			ptrTracker[ptrKey] = true
+			formatAndStoreResult(result.IPAddr, result.PTR, "PTR")
 		}
 	}
+}
 
-	for _, output := range outputs {
-		if utils.OptionCmd.Silent {
-			gologger.Silent().Msgf(output)
-		} else {
-			gologger.Info().Msgf(output)
+// Checks if the result matches the regex for both domain and PTR
+func matchResultWithRegex(result utils.Result, regex string) (bool, bool) {
+	if regex == "" {
+		return true, true
+	}
+	re := regexp.MustCompile(regex)
+	return re.MatchString(result.Domain), re.MatchString(result.PTR)
+}
+
+// Filters the result based on the VirusTotal score
+func filterByVtScore(result utils.Result) bool {
+	virustotal := utils.Virustotal{}
+	result.VtScore = virustotal.GetVtReport(result.Domain)
+	score, err := strconv.ParseUint(utils.OptionCmd.VtscoreValue, 10, 64)
+	if err != nil {
+		gologger.Fatal().Msg("Invalid value for vt-score")
+		return false
+	}
+	return result.VtScore >= score
+}
+
+// Formats and stores the result in the outputs map
+func formatAndStoreResult(ip, value, resultType string) {
+	formattedValue := fmt.Sprintf("%s: %s", resultType, value)
+	if utils.OptionCmd.Silent {
+		outputs[ip] = append(outputs[ip], value)
+	} else {
+		prefix := "├─"
+		if resultType == "PTR" {
+			prefix = "└─"
 		}
+		outputs[ip] = append(outputs[ip], fmt.Sprintf("%s %s", prefix, formattedValue))
 	}
 }
 
-func prepareOutput(result utils.Result, bMatchedDomain bool, bMatchedPTR bool) string {
-	var output = ""
+// Displays the formatted outputs with colors and without PTR duplication
+func displayOutputs() {
+	ipColor := color.New(color.FgCyan).SprintFunc()
+	domainColor := color.New(color.FgGreen).SprintFunc() // Green for domains
+	ptrColor := color.New(color.FgYellow).SprintFunc()   // Yellow for PTRs
 
-	if bMatchedDomain && result.Domain != "" {
-		if utils.OptionCmd.Silent {
-			output = fmt.Sprintf("%s\n", result.Domain)
-		} else {
-			output = fmt.Sprintf("[%s] domain: %s", result.IPAddr, result.Domain)
+	for ip, entries := range outputs {
+		gologger.Print().Msg("──────────────────────────────────────────")
+		gologger.Print().Msgf(" IP: %s", ipColor(ip))
+
+		var domains, ptrs []string
+		for _, entry := range entries {
+			if strings.Contains(entry, "Domain:") {
+				domains = append(domains, entry)
+			} else {
+				ptrs = append(ptrs, entry)
+			}
 		}
+
+		displayEntries(domains, domainColor)
+		displayEntries(ptrs, ptrColor)
 	}
+}
 
-	if bMatchedPTR && result.PTR != "" && output == "" {
+// Displays the entries formatted with colors
+func displayEntries(entries []string, colorFunc func(a ...interface{}) string) {
+	for _, entry := range entries {
+		coloredEntry := colorizeEntry(entry, colorFunc)
 		if utils.OptionCmd.Silent {
-			output = fmt.Sprintf("%s\n", result.PTR)
+			gologger.Silent().Msgf(coloredEntry)
 		} else {
-			output = fmt.Sprintf("[%s] PTR: %s", result.IPAddr, result.PTR)
+			gologger.Print().Msgf(" %s", coloredEntry)
 		}
 	}
+}
 
-	if !utils.OptionCmd.Silent {
-		if utils.OptionCmd.Vtscore {
-			output = fmt.Sprintf("%s VT Score: %d", output, result.VtScore)
-		}
+// Applies color to the domain or PTR value
+func colorizeEntry(entry string, colorFunc func(a ...interface{}) string) string {
+	re := regexp.MustCompile(`: (.+)`)
+	match := re.FindStringSubmatch(entry)
+	if len(match) > 1 {
+		coloredValue := colorFunc(match[1])
+		return re.ReplaceAllString(entry, fmt.Sprintf(": %s", coloredValue))
 	}
-	return output
+	return entry
 }

go.mod

@@ -5,6 +5,7 @@ go 1.21
 require (
 	github.com/PuerkitoBio/goquery v1.9.2
 	github.com/corpix/uarand v0.2.0
+	github.com/fatih/color v1.15.0
 	github.com/projectdiscovery/goflags v0.1.50
 	github.com/projectdiscovery/gologger v1.1.12
 	github.com/projectdiscovery/retryablehttp-go v1.0.58
@@ -32,6 +33,8 @@ require (
 	github.com/klauspost/compress v1.17.8 // indirect
 	github.com/klauspost/pgzip v1.2.6 // indirect
 	github.com/logrusorgru/aurora v2.0.3+incompatible // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/mholt/archiver/v3 v3.5.1 // indirect
 	github.com/microcosm-cc/bluemonday v1.0.26 // indirect
 	github.com/miekg/dns v1.1.59 // indirect

go.sum

@@ -39,6 +39,8 @@ github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDD
 github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5 h1:iFaUwBSo5Svw6L7HYpRu/0lE3e0BaElwnNO1qkNQxBY=
 github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5/go.mod h1:qssHWj60/X5sZFNxpG4HBPDHVqxNm4DfnCKgrbZOT+s=
 github.com/dsnet/golib v0.0.0-20171103203638-1ea166775780/go.mod h1:Lj+Z9rebOhdfkVLjJ8T6VcRQv3SXugXy999NBtR9aFY=
+github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs=
+github.com/fatih/color v1.15.0/go.mod h1:0h5ZqXfHYED7Bhv2ZJamyIOUej9KtShiJESRwBDUSsw=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4HY=
 github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
@@ -85,6 +87,11 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/logrusorgru/aurora v2.0.3+incompatible h1:tOpm7WcpBTn4fjmVfgpQq0EfczGlG91VSDkswnjF5A8=
 github.com/logrusorgru/aurora v2.0.3+incompatible/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mholt/archiver/v3 v3.5.1 h1:rDjOBX9JSF5BvoJGvjqK479aL70qh9DIpZCl+k7Clwo=
 github.com/mholt/archiver/v3 v3.5.1/go.mod h1:e3dqJ7H78uzsRSEACH1joayhuSyhnonssnDhppzS1L4=
 github.com/microcosm-cc/bluemonday v1.0.26 h1:xbqSvqzQMeEHCqMi64VAs4d8uy6Mequs3rQ0k/Khz58=
@@ -256,6 +263,7 @@ golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

utils/banner.go

@@ -1,23 +1,64 @@
 package utils
 
 import (
+	"fmt"
+	"math/rand"
+	"strings"
+	"time"
+
 	"github.com/projectdiscovery/gologger"
 )
 
-var banner = `
+const (
+	primaryColor   = "\033[0;36m" // Cyan
+	secondaryColor = "\033[0;35m" // Magenta
+	reset          = "\033[0m"
+)
 
-╭╮╱╭┳━━━┳━━━╮╱╱╱╱╱╭━━━╮╱╱╭╮╱╱╱╱╱╱╱╱╭╮
-┃┃╱┃┃╭━━┻╮╭╮┃╱╱╱╱╱┃╭━━╯╱╭╯╰╮╱╱╱╱╱╱╭╯╰╮
-┃╰━╯┃╰━━╮┃┃┃┣━╮╭━━┫╰━━┳╮┣╮╭╋━┳━━┳━┻╮╭╋━━┳━╮
-┃╭━╮┃╭━━╯┃┃┃┃╭╮┫━━┫╭━━┻╋╋┫┃┃╭┫╭╮┃╭━┫┃┃╭╮┃╭╯
-┃┃╱┃┃╰━━┳╯╰╯┃┃┃┣━━┃╰━━┳╋╋┫╰┫┃┃╭╮┃╰━┫╰┫╰╯┃┃
-╰╯╱╰┻━━━┻━━━┻╯╰┻━━┻━━━┻╯╰┻━┻╯╰╯╰┻━━┻━┻━━┻╯
-`
+var glitchTitle = []string{
+	"██╗  ██╗███████╗██████╗ ██╗  ██╗",
+	"██║  ██║██╔════╝██╔══██╗╚██╗██╔╝",
+	"███████║█████╗  ██║  ██║ ╚███╔╝ ",
+	"██╔══██║██╔══╝  ██║  ██║ ██╔██╗ ",
+	"██║  ██║███████╗██████╔╝██╔╝ ██╗",
+	"╚═╝  ╚═╝╚══════╝╚═════╝ ╚═╝  ╚═╝",
+}
+
+var subtitle = "HEDnsExtractor - HuntDownProject"
 
 var version = "v1.0.7"
 
 func ShowBanner() {
-	gologger.Print().Msgf("%s\n", banner)
-	gologger.Info().Msgf("Current hednsextractor version %s", version)
-	gologger.Info().Msgf("HEDnsExtractor Config Directory: %s", configDir)
+	displayGlitchTitle()
+	fmt.Println()
+	gologger.Info().Msgf("Current version: %s", version)
+	gologger.Info().Msgf("Config Directory: %s", configDir)
+}
+
+func displayGlitchTitle() {
+	glitchChars := "!@#$%^&*()_+-=[]{}|;:,.<>?"
+	rand.Seed(time.Now().UnixNano())
+
+	for _, line := range glitchTitle {
+		glitchedLine := ""
+		for _, char := range line {
+			if rand.Float32() < 0.1 { // 10% chance of glitch
+				glitchedLine += string(glitchChars[rand.Intn(len(glitchChars))])
+			} else {
+				glitchedLine += string(char)
+			}
+		}
+
+		if rand.Float32() < 0.5 {
+			fmt.Println(primaryColor + glitchedLine + reset)
+		} else {
+			fmt.Println(secondaryColor + glitchedLine + reset)
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	// Display subtitle
+	fmt.Println(primaryColor + strings.Repeat("=", len(subtitle)) + reset)
+	fmt.Println(secondaryColor + subtitle + reset)
+	fmt.Println(primaryColor + strings.Repeat("=", len(subtitle)) + reset)
 }