Merge pull request #337 from tracertea/feature/source_tracking

Added support Multi-source tracking for hosts

v2/pkg/runner/enumerate.go

@@ -44,6 +44,8 @@ func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output strin
 	wg.Add(1)
 	// Create a unique map for filtering duplicate subdomains out
 	uniqueMap := make(map[string]resolve.HostEntry)
+	// Create a map to track sources for each host
+	sourceMap := make(map[string]map[string]struct{})
 	// Process the results in a separate goroutine
 	go func() {
 		for result := range passiveResults {
@@ -57,6 +59,18 @@ func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output strin
 				}
 				subdomain := strings.ReplaceAll(strings.ToLower(result.Value), "*.", "")
 
+				if _, ok := uniqueMap[subdomain]; !ok {
+					sourceMap[subdomain] = make(map[string]struct{})
+
+				}
+
+				// Log the verbose message about the found subdomain per source
+				if _, ok := sourceMap[subdomain][result.Source]; !ok{
+					gologger.Verbosef("%s\n", result.Source, subdomain)
+				}
+
+				sourceMap[subdomain][result.Source] = struct{}{}
+
 				// Check if the subdomain is a duplicate. If not,
 				// send the subdomain for resolution.
 				if _, ok := uniqueMap[subdomain]; ok {
@@ -67,10 +81,6 @@ func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output strin
 
 				uniqueMap[subdomain] = hostEntry
 
-				// Log the verbose message about the found subdomain and send the
-				// host for resolution to the resolution pool
-				gologger.Verbosef("%s\n", result.Source, subdomain)
-
 				// If the user asked to remove wildcard then send on the resolve
 				// queue. Otherwise, if mode is not verbose print the results on
 				// the screen as they are discovered.
@@ -116,7 +126,11 @@ func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output strin
 		if r.options.RemoveWildcard {
 			err = outputter.WriteHostNoWildcard(foundResults, os.Stdout)
 		} else {
-			err = outputter.WriteHost(uniqueMap, os.Stdout)
+			if r.options.CaptureSources {
+				err = outputter.WriteSourceHost(sourceMap,os.Stdout)
+			} else {
+				err = outputter.WriteHost(uniqueMap, os.Stdout)
+			}
 		}
 	}
 	if err != nil {
@@ -167,7 +181,11 @@ func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output strin
 			if r.options.RemoveWildcard {
 				err = outputter.WriteHostNoWildcard(foundResults, file)
 			} else {
-				err = outputter.WriteHost(uniqueMap, file)
+				if r.options.CaptureSources {
+					err = outputter.WriteSourceHost(sourceMap, file)
+				} else {
+					err = outputter.WriteHost(uniqueMap, file)
+				}
 			}
 		}
 		if err != nil {

v2/pkg/runner/options.go

@@ -21,6 +21,7 @@ type Options struct {
 	Silent             bool   // Silent suppresses any extra text and only writes subdomains to screen
 	ListSources        bool   // ListSources specifies whether to list all available sources
 	RemoveWildcard     bool   // RemoveWildcard specifies whether to remove potential wildcard or dead subdomains from the results.
+	CaptureSources     bool   // CaptureSources specifies whether to save all sources that returned a specific domains or just the first source
 	Stdin              bool   // Stdin specifies whether stdin input was given to the process
 	Version            bool   // Version specifies if we should just show version and exit
 	Recursive          bool   // Recursive specifies whether to use only recursive subdomain enumeration sources
@@ -62,6 +63,7 @@ func ParseOptions() *Options {
 	flag.StringVar(&options.Output, "o", "", "File to write output to (optional)")
 	flag.StringVar(&options.OutputDirectory, "oD", "", "Directory to write enumeration results to (optional)")
 	flag.BoolVar(&options.JSON, "json", false, "Write output in JSON lines Format")
+	flag.BoolVar(&options.CaptureSources, "collect-sources", false, "Output host source as array of sources instead of single (first) source")
 	flag.BoolVar(&options.JSON, "oJ", false, "Write output in JSON lines Format")
 	flag.BoolVar(&options.HostIP, "oI", false, "Write output in Host,IP format")
 	flag.BoolVar(&options.Silent, "silent", false, "Show only subdomains in output")

v2/pkg/runner/outputter.go

@@ -23,6 +23,17 @@ type jsonResult struct {
 	Source string `json:"source"`
 }
 
+type jsonSourceResultIP struct {
+	Host   string `json:"host"`
+	IP     string `json:"ip"`
+	Sources []string `json:"sources"`
+}
+
+type jsonSourceResult struct {
+	Host   string `json:"host"`
+	Sources []string `json:"sources"`
+}
+
 // NewOutputter creates a new Outputter
 func NewOutputter(json bool) *OutPutter {
 	return &OutPutter{JSON: json}
@@ -178,3 +189,59 @@ func writeJSONHost(results map[string]resolve.HostEntry, writer io.Writer) error
 	}
 	return nil
 }
+
+// WriteHostIP writes the output list of subdomain to an io.Writer
+func (o *OutPutter) WriteSourceHost(sourceMap map[string]map[string]struct{}, writer io.Writer) error {
+	var err error
+	if o.JSON {
+		err = writeSourceJSONHost(sourceMap, writer)
+	} else {
+		err = writeSourcePlainHost(sourceMap, writer)
+	}
+	return err
+}
+
+func writeSourceJSONHost(sourceMap map[string]map[string]struct{}, writer io.Writer) error {
+	encoder := jsoniter.NewEncoder(writer)
+
+	var data jsonSourceResult
+
+	for host, sources := range sourceMap {
+		data.Host = host
+		keys := make([]string, 0, len(sources))
+		for source := range sources {
+			keys = append(keys, source)
+		}
+		data.Sources = keys
+
+		err := encoder.Encode(&data)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func writeSourcePlainHost(sourceMap map[string]map[string]struct{}, writer io.Writer) error {
+	bufwriter := bufio.NewWriter(writer)
+	sb := &strings.Builder{}
+
+	for host, sources := range sourceMap {
+		sb.WriteString(host)
+		sb.WriteString(",[")
+		sourcesString := ""
+		for source := range sources {
+			sourcesString += source + ","
+		}
+		sb.WriteString(strings.Trim(sourcesString, ", "))
+		sb.WriteString("]\n")
+
+		_, err := bufwriter.WriteString(sb.String())
+		if err != nil {
+			bufwriter.Flush()
+			return err
+		}
+		sb.Reset()
+	}
+	return bufwriter.Flush()
+}