Added logger, hopefully final

cmd/shuffledns/main.go

@@ -9,9 +9,11 @@ func main() {
 	// Parse the command line flags and read config files
 	options := runner.ParseOptions()
 
-	runner, err := runner.NewRunner(options)
+	runner, err := runner.New(options)
 	if err != nil {
 		gologger.Fatalf("Could not create runner: %s\n", err)
 	}
 
+	runner.RunEnumeration()
+	runner.Close()
 }

pkg/massdns/process.go

@@ -9,7 +9,9 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"time"
 
+	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/shuffledns/internal/store"
 	"github.com/projectdiscovery/shuffledns/pkg/parser"
 	"github.com/rs/xid"
@@ -28,23 +30,31 @@ func (c *Client) Process() error {
 
 	// Create a store for storing ip metadata
 	store := store.New()
+	defer store.Close()
 
 	// Create a temporary file for the massdns output
 	temporaryOutput := path.Join(c.config.TempDir, xid.New().String())
 
+	gologger.Infof("Creating temporary massdns output file: %s\n", temporaryOutput)
+	gologger.Infof("Executing massdns on %s\n", c.config.Domain)
+
+	now := time.Now()
 	// Run the command on a temp file and wait for the output
 	cmd := exec.Command(c.config.MassdnsPath, []string{"-r", c.config.ResolversFile, "-t", "A", c.config.InputFile, "-w", temporaryOutput, "-s", strconv.Itoa(c.config.Threads)}...)
 	err = cmd.Run()
 	if err != nil {
 		return fmt.Errorf("could not execute massdns: %w", err)
 	}
+	gologger.Infof("Massdns execution took %s\n", time.Now().Sub(now))
 
 	massdnsOutput, err := os.Open(temporaryOutput)
 	if err != nil {
 		return fmt.Errorf("could not open massdns output file: %w", err)
 	}
 	defer massdnsOutput.Close()
 
+	gologger.Infof("Parsing output and removing wildcards\n")
+
 	err = parser.Parse(massdnsOutput, func(domain string, ip []string) {
 		for _, ip := range ip {
 			// We've stumbled upon a wildcard, just ignore it.
@@ -85,19 +95,33 @@ func (c *Client) Process() error {
 				}
 			}
 			// Put the new hostname and increment the counter by 1.
-			record.Hostnames = append(record.Hostnames, resp.Domain)
+			record.Hostnames = append(record.Hostnames, domain)
 			record.Counter++
 		}
 	})
 	if err != nil {
 		return fmt.Errorf("could not parse massdns output: %w", err)
 	}
+	gologger.Infof("Finished enumeration, started writing output\n")
 
-	output, err := os.Create(c.config.OutputFile)
-	if err != nil {
-		return fmt.Errorf("could not create massdns output file: %v", err)
+	// Parse the massdns output
+	return c.writeOutput(store)
+}
+
+func (c *Client) writeOutput(store *store.Store) error {
+	// Write the unique deduplicated output to the file or stdout
+	// depending on what the user has asked.
+	var output *os.File
+	var w *bufio.Writer
+	var err error
+
+	if c.config.OutputFile != "" {
+		output, err = os.Create(c.config.OutputFile)
+		if err != nil {
+			return fmt.Errorf("could not create massdns output file: %v", err)
+		}
+		w = bufio.NewWriter(output)
 	}
-	w := bufio.NewWriter(output)
 	buffer := &strings.Builder{}
 
 	uniqueMap := make(map[string]struct{})
@@ -110,14 +134,20 @@ func (c *Client) Process() error {
 			// Check if we don't have a duplicate
 			if _, ok := uniqueMap[data]; !ok {
 				uniqueMap[data] = struct{}{}
-				w.WriteString(data)
+
+				if output != nil {
+					w.WriteString(data)
+				}
+				gologger.Silentf("%s", data)
 				buffer.Reset()
 			}
 		}
 	}
-	w.Flush()
-	output.Close()
-	store.Close()
 
+	// Close the files and return
+	if output != nil {
+		w.Flush()
+		output.Close()
+	}
 	return nil
 }

pkg/runner/options.go

@@ -71,7 +71,7 @@ func ParseOptions() *Options {
 
 	// Set the domain in the config if provided by user from the stdin
 	if options.Stdin && options.Wordlist != "" {
-		buffer := bytes.Buffer{}
+		buffer := &bytes.Buffer{}
 		io.Copy(buffer, os.Stdin)
 		options.Domain = buffer.String()
 	}

pkg/runner/runner.go

@@ -7,8 +7,10 @@ import (
 	"io/ioutil"
 	"os"
 	"path"
+	"time"
 
 	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/shuffledns/pkg/massdns"
 	"github.com/rs/xid"
 )
 
@@ -31,6 +33,7 @@ func New(options *Options) (*Runner, error) {
 		if options.MassdnsPath == "" {
 			return nil, errors.New("could not find massdns binary")
 		}
+		gologger.Debugf("Discovered massdns binary at %s\n", options.MassdnsPath)
 	}
 
 	// Create a temporary directory that will be removed at the end
@@ -65,48 +68,55 @@ func (r *Runner) findBinary() string {
 	return ""
 }
 
-// runEnumeration sets up the input layer for giving input to massdns
+// RunEnumeration sets up the input layer for giving input to massdns
 // binary and runs the actual enumeration
-func (r *Runner) runEnumeration() error {
+func (r *Runner) RunEnumeration() {
 	// Handle stdin input
 	if r.options.Stdin {
 		// Is the stdin input a domain for bruteforce
-		if r.options.Wordlist {
+		if r.options.Wordlist != "" {
 			r.processDomain()
+			return
 		}
 		// Write the input from stdin to a file and resolve it.
 		r.processSubdomains()
+		return
 	}
 
 	// Handle a list of subdomains to resolve
-	if r.options.SubdomainsList {
+	if r.options.SubdomainsList != "" {
 		r.processSubdomains()
+		return
 	}
 
 	// Handle a domain to bruteforce with wordlist
-	if r.options.Wordlist {
+	if r.options.Wordlist != "" {
 		r.processDomain()
+		return
 	}
 }
 
 // processDomain processes the bruteforce for a domain using a wordlist
 func (r *Runner) processDomain() {
 	resolveFile := path.Join(r.tempDir, xid.New().String())
-	file, err := os.Create(resolverFile)
+	file, err := os.Create(resolveFile)
 	if err != nil {
-		gologger.Fatalf("Could not create bruteforce list (%s): %s\n", r.tempDir, err)
+		gologger.Errorf("Could not create bruteforce list (%s): %s\n", r.tempDir, err)
 		return
 	}
 	writer := bufio.NewWriter(file)
 
 	// Read the input wordlist for bruteforce generation
 	inputFile, err := os.Open(r.options.Wordlist)
 	if err != nil {
-		gologger.Fatalf("Could not read bruteforce wordlist (%s): %s\n", r.options.Wordlist, err)
+		gologger.Errorf("Could not read bruteforce wordlist (%s): %s\n", r.options.Wordlist, err)
 		file.Close()
 		return
 	}
 
+	gologger.Infof("Started generating bruteforce permutation\n")
+
+	now := time.Now()
 	// Create permutation for domain with wordlist
 	scanner := bufio.NewScanner(inputFile)
 	for scanner.Scan() {
@@ -120,6 +130,10 @@ func (r *Runner) processDomain() {
 	inputFile.Close()
 	file.Close()
 
+	gologger.Infof("Generating permutations took %s\n", time.Now().Sub(now))
+
+	// Run the actual massdns enumeration process
+	r.runMassdns(resolveFile)
 }
 
 // processSubdomain processes the resolving for a list of subdomains
@@ -129,9 +143,9 @@ func (r *Runner) processSubdomains() {
 	// If there is stdin, write the resolution list to the file
 	if r.options.Stdin {
 		resolveFile = path.Join(r.tempDir, xid.New().String())
-		file, err := os.Create(resolverFile)
+		file, err := os.Create(resolveFile)
 		if err != nil {
-			gologger.Fatalf("Could not create resolution list (%s): %s\n", tempDir, err)
+			gologger.Errorf("Could not create resolution list (%s): %s\n", r.tempDir, err)
 			return
 		}
 		io.Copy(file, os.Stdin)
@@ -140,4 +154,31 @@ func (r *Runner) processSubdomains() {
 		// Use the file if user has provided one
 		resolveFile = r.options.SubdomainsList
 	}
+
+	// Run the actual massdns enumeration process
+	r.runMassdns(resolveFile)
+}
+
+// runMassdns runs the massdns tool on the list of inputs
+func (r *Runner) runMassdns(inputFile string) {
+	massdns, err := massdns.New(massdns.Config{
+		Domain:        r.options.Domain,
+		Retries:       r.options.Retries,
+		MassdnsPath:   r.options.MassdnsPath,
+		Threads:       r.options.Threads,
+		InputFile:     inputFile,
+		ResolversFile: r.options.ResolversFile,
+		TempDir:       r.tempDir,
+		OutputFile:    r.options.Output,
+	})
+	if err != nil {
+		gologger.Errorf("Could not create massdns client: %s\n", err)
+		return
+	}
+
+	err = massdns.Process()
+	if err != nil {
+		gologger.Errorf("Could not run massdns: %s\n", err)
+	}
+	gologger.Infof("Finished resolving. Hack the Planet!\n")
 }