Skip to content

Commit

Permalink
Merge pull request #413 from hachh/go-library
Browse files Browse the repository at this point in the history 
make EnumerateSingleDomain usable as a Go library
  • Loading branch information
@ehsandeep
ehsandeep authored Apr 17, 2021
2 parents 9e7ee80 + 102fdd5 commit 1963910
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 54 deletions.
61 changes: 17 additions & 44 deletions v2/pkg/runner/enumerate.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ package runner

import (
"context"
"os"
"io"
"strings"
"sync"
"time"
Expand All @@ -16,7 +16,7 @@ import (
const maxNumCount = 2

// EnumerateSingleDomain performs subdomain enumeration against a single domain
func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output string, appendToFile bool) error {
func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain string, outputs []io.Writer) error {
gologger.Info().Msgf("Enumerating subdomains for %s\n", domain)

// Get the API keys for sources from the configuration
Expand Down Expand Up @@ -115,62 +115,35 @@ func (r *Runner) EnumerateSingleDomain(ctx context.Context, domain, output strin

outputter := NewOutputter(r.options.JSON)

// If verbose mode was used, then now print all the
// found subdomains on the screen together.
// Now output all results in output writers
var err error
if r.options.HostIP {
err = outputter.WriteHostIP(foundResults, os.Stdout)
} else {
if r.options.RemoveWildcard {
err = outputter.WriteHostNoWildcard(foundResults, os.Stdout)
} else {
if r.options.CaptureSources {
err = outputter.WriteSourceHost(sourceMap, os.Stdout)
} else {
err = outputter.WriteHost(uniqueMap, os.Stdout)
}
}
}
if err != nil {
gologger.Error().Msgf("Could not verbose results for %s: %s\n", domain, err)
return err
}

// Show found subdomain count in any case.
duration := durafmt.Parse(time.Since(now)).LimitFirstN(maxNumCount).String()
if r.options.RemoveWildcard {
gologger.Info().Msgf("Found %d subdomains for %s in %s\n", len(foundResults), domain, duration)
} else {
gologger.Info().Msgf("Found %d subdomains for %s in %s\n", len(uniqueMap), domain, duration)
}

if output != "" {
file, err := outputter.createFile(output, appendToFile)
if err != nil {
gologger.Error().Msgf("Could not create file %s for %s: %s\n", output, domain, err)
return err
}

defer file.Close()

for _, w := range outputs {
if r.options.HostIP {
err = outputter.WriteHostIP(foundResults, file)
err = outputter.WriteHostIP(foundResults, w)
} else {
if r.options.RemoveWildcard {
err = outputter.WriteHostNoWildcard(foundResults, file)
err = outputter.WriteHostNoWildcard(foundResults, w)
} else {
if r.options.CaptureSources {
err = outputter.WriteSourceHost(sourceMap, file)
err = outputter.WriteSourceHost(sourceMap, w)
} else {
err = outputter.WriteHost(uniqueMap, file)
err = outputter.WriteHost(uniqueMap, w)
}
}
}
if err != nil {
gologger.Error().Msgf("Could not write results to file %s for %s: %s\n", output, domain, err)
gologger.Error().Msgf("Could not verbose results for %s: %s\n", domain, err)
return err
}
}

// Show found subdomain count in any case.
duration := durafmt.Parse(time.Since(now)).LimitFirstN(maxNumCount).String()
if r.options.RemoveWildcard {
gologger.Info().Msgf("Found %d subdomains for %s in %s\n", len(foundResults), domain, duration)
} else {
gologger.Info().Msgf("Found %d subdomains for %s in %s\n", len(uniqueMap), domain, duration)
}

return nil
}
9 changes: 7 additions & 2 deletions v2/pkg/runner/options.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package runner

import (
"flag"
"io"
"os"
"path"
"reflect"
Expand Down Expand Up @@ -30,7 +31,8 @@ type Options struct {
MaxEnumerationTime int // MaxEnumerationTime is the maximum amount of time in mins to wait for enumeration
Domain string // Domain is the domain to find subdomains for
DomainsFile string // DomainsFile is the file containing list of domains to find subdomains for
Output string // Output is the file to write found subdomains to.
Output io.Writer
OutputFile string // Output is the file to write found subdomains to.
OutputDirectory string // OutputDirectory is the directory to write results to in case list of domains is given
Sources string // Sources contains a comma-separated list of sources to use for enumeration
ExcludeSources string // ExcludeSources contains the comma-separated sources to not include in the enumeration process
Expand Down Expand Up @@ -58,7 +60,7 @@ func ParseOptions() *Options {
flag.IntVar(&options.MaxEnumerationTime, "max-time", 10, "Minutes to wait for enumeration results")
flag.StringVar(&options.Domain, "d", "", "Domain to find subdomains for")
flag.StringVar(&options.DomainsFile, "dL", "", "File containing list of domains to enumerate")
flag.StringVar(&options.Output, "o", "", "File to write output to (optional)")
flag.StringVar(&options.OutputFile, "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")
Expand All @@ -77,6 +79,9 @@ func ParseOptions() *Options {
flag.BoolVar(&options.Version, "version", false, "Show version of subfinder")
flag.Parse()

// Default output is stdout
options.Output = os.Stdout

// Check if stdin pipe was given
options.Stdin = hasStdin()

Expand Down
52 changes: 44 additions & 8 deletions v2/pkg/runner/runner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"os"
"path"

"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/subfinder/v2/pkg/passive"
"github.com/projectdiscovery/subfinder/v2/pkg/resolve"
)
Expand Down Expand Up @@ -39,9 +40,24 @@ func NewRunner(options *Options) (*Runner, error) {

// RunEnumeration runs the subdomain enumeration flow on the targets specified
func (r *Runner) RunEnumeration(ctx context.Context) error {
outputs := []io.Writer{r.options.Output}

// Check if only a single domain is sent as input. Process the domain now.
if r.options.Domain != "" {
return r.EnumerateSingleDomain(ctx, r.options.Domain, r.options.Output, false)
// If output file specified, create file
if r.options.OutputFile != "" {
outputter := NewOutputter(r.options.JSON)
file, err := outputter.createFile(r.options.OutputFile, false)
if err != nil {
gologger.Error().Msgf("Could not create file %s for %s: %s\n", r.options.OutputFile, r.options.Domain, err)
return err
}
defer file.Close()

outputs = append(outputs, file)
}

return r.EnumerateSingleDomain(ctx, r.options.Domain, outputs)
}

// If we have multiple domains as input,
Expand All @@ -50,21 +66,21 @@ func (r *Runner) RunEnumeration(ctx context.Context) error {
if err != nil {
return err
}
err = r.EnumerateMultipleDomains(ctx, f)
err = r.EnumerateMultipleDomains(ctx, f, outputs)
f.Close()
return err
}

// If we have STDIN input, treat it as multiple domains
if r.options.Stdin {
return r.EnumerateMultipleDomains(ctx, os.Stdin)
return r.EnumerateMultipleDomains(ctx, os.Stdin, outputs)
}
return nil
}

// EnumerateMultipleDomains enumerates subdomains for multiple domains
// We keep enumerating subdomains for a given domain until we reach an error
func (r *Runner) EnumerateMultipleDomains(ctx context.Context, reader io.Reader) error {
func (r *Runner) EnumerateMultipleDomains(ctx context.Context, reader io.Reader, outputs []io.Writer) error {
scanner := bufio.NewScanner(reader)
for scanner.Scan() {
domain := scanner.Text()
Expand All @@ -73,21 +89,41 @@ func (r *Runner) EnumerateMultipleDomains(ctx context.Context, reader io.Reader)
}

var err error
var file *os.File
// If the user has specified an output file, use that output file instead
// of creating a new output file for each domain. Else create a new file
// for each domain in the directory.
if r.options.Output != "" {
err = r.EnumerateSingleDomain(ctx, domain, r.options.Output, true)
if r.options.OutputFile != "" {
outputter := NewOutputter(r.options.JSON)
file, err = outputter.createFile(r.options.OutputFile, true)
if err != nil {
gologger.Error().Msgf("Could not create file %s for %s: %s\n", r.options.OutputFile, r.options.Domain, err)
return err
}

err = r.EnumerateSingleDomain(ctx, domain, append(outputs, file))

file.Close()
} else if r.options.OutputDirectory != "" {
outputFile := path.Join(r.options.OutputDirectory, domain)
if r.options.JSON {
outputFile += ".json"
} else {
outputFile += ".txt"
}
err = r.EnumerateSingleDomain(ctx, domain, outputFile, false)

outputter := NewOutputter(r.options.JSON)
file, err = outputter.createFile(outputFile, false)
if err != nil {
gologger.Error().Msgf("Could not create file %s for %s: %s\n", r.options.OutputFile, r.options.Domain, err)
return err
}

err = r.EnumerateSingleDomain(ctx, domain, append(outputs, file))

file.Close()
} else {
err = r.EnumerateSingleDomain(ctx, domain, "", true)
err = r.EnumerateSingleDomain(ctx, domain, outputs)
}
if err != nil {
return err
Expand Down

0 comments on commit 1963910

Please sign in to comment.